diff --git a/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/README.md b/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/README.md index 36f7d4726..e0410cbbf 100644 --- a/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/README.md +++ b/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/README.md @@ -36,6 +36,6 @@ Canonical benchmark helpers must be materialized by Arena; do not edit their gen ## Protected evaluation controls -Controls cover remapped requests, a prefill early return and heterogeneous numbers of logits. Both mutated input IDs and returned int64 logit indices are checked. +Controls cover a prefill early return and remapped decode requests with zero, one or two draft tokens. A further correctness-only control builds draft-token rows with padding on the target device, so the wrapper receives the actual row stride. Mutated input IDs and returned int64 logit indices must match the exact reference; draft tokens and their padding must remain unchanged. The task-local `_arena_contract.py` and `_arena_replay.py` are protected evaluation code. Original cases, seeds, tolerances, warmups, sample counts, allocations and preparation boundaries remain in `scripts/task_runner.py`. The extra `contract_controls` manifest row is correctness-only. Both the frozen baseline and candidate receive the same checks. The measured graph exposes its real outputs; an untimed replay changes a domain-valid input, recomputes the CPU oracle and restores all input buffers in `finally`. For the zero operator the replay control instead poisons its output. References and snapshots are outside device timing. Failure to observe or replay the measured invocation is an error, never an accepted timing sample. diff --git a/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/_arena_contract.py b/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/_arena_contract.py index 313f04ba2..638358103 100644 --- a/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/_arena_contract.py +++ b/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/_arena_contract.py @@ -16,6 +16,26 @@ def controls(harness, function, device): for args in control_inputs(harness): function(*to_device(args, device)) + args, storage = padded_draft_inputs(harness, device) + draft_tokens = args[6] + assert draft_tokens.stride() == (5, 1) + drafts_before = draft_tokens.clone() + padding_before = storage[:, 2:].clone() + function(*args) + torch.testing.assert_close(draft_tokens, drafts_before, atol=0, rtol=0) + torch.testing.assert_close(storage[:, 2:], padding_before, atol=0, rtol=0) + + +def padded_draft_inputs(harness, device): + from _arena_replay import to_device + # Construct the view on the target device; moving a CPU slice could erase + # its padded row stride. Reuse the remapped 0/1/2-draft decode inputs. + args = list(to_device(tuple(control_inputs(harness))[1], device)) + storage = torch.full((3, 5), -777, dtype=torch.int32, device=device) + storage[:, :2].copy_(args[6]) + args[6] = storage[:, :2] + return tuple(args), storage + FUNCTION = 'combine_sampled_and_draft_tokens' MUTABLE = (0,) ATOL = 0.0 @@ -29,8 +49,14 @@ def fresh(args): args[2].add_(1); args[6].add_(2) def control_inputs(harness): + # Mixed prefill/decode requests with different numbers of draft tokens. yield (i([-9]*7), i([3,0,2]), l([10,11,12,13]), i([0,1,4,7]), i([3,8,7]), i([4,0,2,3]), i([[1,2,3],[4,5,6],[7,8,9],[10,11,12]]), i([0,1,3,6]), 6) + # The first decode request has no drafts. The other decode requests use + # remapped state rows and contribute two and one drafts, respectively. + yield (i([-9]*6), i([2,0,1]), l([101,202,303]), i([0,1,4,6]), + i([10,20,30]), i([0,0,0]), i([[111,112],[211,212],[311,312]]), + i([0,1,4,6]), 6) def observe(result, args): return args[0], result diff --git a/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/workloads.json b/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/workloads.json index 7d6b36185..fa6d9a602 100644 --- a/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/workloads.json +++ b/tasks/triton2triton/vllm/triton_combine_sampled_and_draft_tokens/workloads.json @@ -135,7 +135,7 @@ ], "params": { "case_index": 5, - "coverage": "Controls cover remapped requests, a prefill early return and heterogeneous numbers of logits. Both mutated input IDs and returned int64 logit indices are checked." + "coverage": "Correctness-only controls cover a prefill early return, remapped decode requests with zero, one or two drafts, and a padded draft row stride. Mutated input IDs, returned int64 logit indices, draft input and padding are checked." } } ], diff --git a/tasks/triton2triton/vllm/triton_pack_bitmatrix/BOUNDARY_CHECKS.md b/tasks/triton2triton/vllm/triton_pack_bitmatrix/BOUNDARY_CHECKS.md new file mode 100644 index 000000000..4a5e30b2a --- /dev/null +++ b/tasks/triton2triton/vllm/triton_pack_bitmatrix/BOUNDARY_CHECKS.md @@ -0,0 +1,50 @@ +# Top-k boundary checks + +The packing kernel must include every selected expert in every row, including +assignments after slot 31 and in intermediate tiles. For valid expert IDs in +`[0, num_experts)`, set bit `id % 32` in output word `id // 32`. Duplicate IDs +set a single bit. The output has shape `(rows, ceil(num_experts / 32))` and dtype +`uint32` on the input device. Masked tile-padding slots must not set any bits. +Valid two-dimensional `int16` inputs may have row padding, column gaps, or +both. Packing uses their logical values and leaves the input unchanged. + +The first three `control-upstream-*` cases retain the PR105 controls for +duplicates, expert-word boundaries and row tails. The next two cases are +repository-authored top-k controls, identified by their `source` metadata in +`workloads.json`; their existing case IDs are retained for compatibility. +They use multiple rows, top-k 33 and 65, and 70 experts spanning three output +words. Each tile contributes expert bits that cannot be supplied by the other +tiles, and rows have different memberships. The native correctness action +rejects candidates that skip later tiles, omit the middle tile, or read another +row's later assignments. Partial final tiles and output-word tails are checked. +Three more repository-authored native controls use noncontiguous GPU inputs +with shape `(3, 2)` and 34 experts. Their row, column and combined-gap strides +make the old row-major linear load return incorrect expert bits; the controls +compare against the same independent membership oracle. + +The original five scored shapes, exact-equality checks, 10 warmups and 100 +samples remain unchanged. Graph timing can batch multiple wrapper calls in one +replay and report time per call. The benchmark guard checks every output from +the final captured batch, including after poisoning every output and replaying +the same graph with changed input IDs. Outputs from outstanding calls must not +touch the same bytes; disjoint views of one allocation, including interleaved +views, are allowed. This prevents +a wrapper from doing work only for selected calls while returning one shared +result for the others. Event timing still checks its one measured call and replay. +These output checks do not prove that every possible cached or copied +implementation performs independent work for identical inputs. + +The implementation keeps the original single-tile branch for the existing +scored workloads. The wrapper makes only noncontiguous inputs contiguous +before that kernel reads them; contiguous scored inputs need no copy. +Simplifying the single-tile branch requires a new paired timing check; +equivalent expressions can produce different generated code. + +This task freezes the initial candidate as its performance baseline. The +top-k packing repair changes that baseline's source version. Record the task +revision and baseline fingerprint when comparing runs; results from different +baseline versions must not be pooled as measurements against one baseline. + +The existing implementation is adapted from the Apache-2.0-licensed +[vLLM packing kernel](https://github.com/vllm-project/vllm/blob/v0.11.0/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py). +Copyright contributors to the vLLM project. diff --git a/tasks/triton2triton/vllm/triton_pack_bitmatrix/_arena_checks.py b/tasks/triton2triton/vllm/triton_pack_bitmatrix/_arena_checks.py index 491be23f3..2a2a981c7 100644 --- a/tasks/triton2triton/vllm/triton_pack_bitmatrix/_arena_checks.py +++ b/tasks/triton2triton/vllm/triton_pack_bitmatrix/_arena_checks.py @@ -1,6 +1,7 @@ """Check true expert membership and the exact measured bitmatrix replay.""" from contextlib import contextmanager import inspect +from itertools import product def check_output(output, expected): @@ -26,6 +27,52 @@ def perturb(ids, experts): ids.copy_((ids + 1).remainder(experts)) +def output_span(output): + """Conservative byte span touched by a tensor, including a strided view.""" + base = output.untyped_storage().data_ptr() + first = last = output.storage_offset() + for size, stride in zip(output.shape, output.stride()): + extent = (size - 1) * stride + first += min(0, extent) + last += max(0, extent) + element_size = output.element_size() + return base + first * element_size, base + (last + 1) * element_size + + +def check_disjoint_outputs(outputs): + # Most captured outputs have non-overlapping bounding spans. Only groups + # whose spans intersect need the exact, stride-aware byte check below. + spans = sorted(((*output_span(output), index) for index, output in enumerate(outputs))) + + def check_group(group): + if len(group) < 2: + return + owners = {} + for _, _, index in group: + output = outputs[index] + base = output.untyped_storage().data_ptr() + element_size = output.element_size() + for position in product(*(range(size) for size in output.shape)): + element = output.storage_offset() + sum( + coordinate * stride for coordinate, stride in zip(position, output.stride())) + first_byte = base + element * element_size + for address in range(first_byte, first_byte + element_size): + previous = owners.setdefault(address, index) + if previous != index: + raise AssertionError( + f'Captured bitmatrix outputs {previous} and {index} overlap') + + group = [] + group_end = 0 + for span in spans: + if group and span[0] >= group_end: + check_group(group) + group = [] + group.append(span) + group_end = max(group_end, span[1]) + check_group(group) + + @contextmanager def checked_modules(harness): load_original = harness.load_module @@ -70,10 +117,12 @@ def checked_benchmark(harness, benchmark, fn, **kwargs): expected = reference(harness, pristine, experts) original = module.pack_topk_to_bitmatrix captured = None + returned = [] def collect(*args): nonlocal captured captured = original(*args) + returned.append(captured) return captured def measured(): @@ -88,17 +137,44 @@ def measured(): try: timed = harness._TimedRun() ms, metadata = benchmark(measured, timed_run=timed, **kwargs) + method = metadata.get('benchmark_method') + kind = metadata.get('benchmark_timed_run_kind') + repeats = metadata.get('benchmark_effective_repeats') + if type(repeats) is not int or repeats < 1: + raise AssertionError('Benchmark omitted a valid effective repeat count') + if method == 'cuda_graph': + if kind != 'captured_graph' or len(returned) < repeats: + raise AssertionError('Final captured graph invocations are not observable') + outputs = returned[-repeats:] + elif method == 'cuda_event_fallback': + if kind != 'eager_callable' or repeats != 1 or not returned: + raise AssertionError('Event timing invocation is not observable') + outputs = returned[-1:] + else: + raise AssertionError('Unrecognized device benchmark method') + if outputs[-1] is not timed.outputs: + raise AssertionError('Measured output is not the final captured invocation') unchanged(ids, pristine) - check_output(timed.outputs, expected) + for output in outputs: + check_output(output, expected) + check_disjoint_outputs(outputs) perturb(ids, experts) replay_pristine = ids.clone() replay_expected = reference(harness, replay_pristine, experts) - timed.outputs.view(torch.int32).fill_(-1) + for output in outputs: + output.fill_(0xFFFFFFFF) replayed = timed.rerun() unchanged(ids, replay_pristine) + if method == 'cuda_graph' and replayed is not outputs[-1]: + raise AssertionError('Replay did not return the final captured output') check_output(replayed, replay_expected) + if method == 'cuda_graph': + for output in outputs: + check_output(output, replay_expected) return ms, {**metadata, 'timed_output_checked': True, - 'perturbed_input_replay_checked': True, 'source_buffers_unchanged': True} + 'perturbed_input_replay_checked': True, 'source_buffers_unchanged': True, + 'timed_invocation_outputs_checked': len(outputs), + 'timed_invocation_outputs_disjoint': True} finally: module.pack_topk_to_bitmatrix = original ids.copy_(pristine) diff --git a/tasks/triton2triton/vllm/triton_pack_bitmatrix/_upstream_controls.py b/tasks/triton2triton/vllm/triton_pack_bitmatrix/_upstream_controls.py index 0b0e5fdf2..02b34bae7 100644 --- a/tasks/triton2triton/vllm/triton_pack_bitmatrix/_upstream_controls.py +++ b/tasks/triton2triton/vllm/triton_pack_bitmatrix/_upstream_controls.py @@ -1,11 +1,13 @@ -"""Additional unscored public-branch controls from PR105. +"""Unscored controls for expert-bit packing. -Source revision: 0acf65b3a967ef1025dbfc5fd4b415b259e3bd43 -Source path: tasks/triton2triton/vllm/triton_pack_bitmatrix/scripts/task_runner.py +The first three controls come from PR105 (revision +0acf65b3a967ef1025dbfc5fd4b415b259e3bd43, task runner path +tasks/triton2triton/vllm/triton_pack_bitmatrix/scripts/task_runner.py). +The final two top-k boundary controls and three strided-input controls were +added in this repository. Original scored inputs, gates and timing remain in scripts/task_runner.py. """ import sys, os, json, argparse, importlib.util -BOUNDARY_TEST_SHAPES = [(17, 31, 1), (33, 32, 31), (513, 33, 32)] def reference_pack_bitmatrix(topk_ids, num_experts): """Independent CPU expert-membership oracle; padding is not an assignment.""" @@ -30,8 +32,46 @@ def make_boundary_topk_ids(n_rows, num_experts, topk, device): prefix_len = min(topk, len(edge_ids)) topk_ids[0, :prefix_len] = torch.tensor(edge_ids[:prefix_len], dtype=torch.int16) topk_ids[-1, 0] = num_experts - 1 + if topk > 32: + # Each row has a distinct first-tile signature and a duplicated ID. + # Later tiles add bits absent from that row's preceding tiles, so + # skipping a tile or reusing row 0's later assignments is observable. + first = [0, 31, 32, 63, min(64, num_experts - 1)] + middle = [33, 34, 35, 36, 37] + final = [num_experts - 1 - i for i in range(5)] + for row in range(n_rows): + slot = row % 5 + topk_ids[row, :32] = first[slot] + topk_ids[row, 1] = first[slot] # An explicit duplicate. + topk_ids[row, 31] = slot + 1 + if topk > 33: + topk_ids[row, 32:64] = middle[slot] + topk_ids[row, -1] = final[slot] return topk_ids.to(device) -EXTRA_CASES = [(17, 31, 1), (33, 32, 31), (513, 33, 32)] + + +def make_strided_topk_ids(n_rows, num_experts, topk, layout, device): + """Keep the physical gaps on the target device for layout controls.""" + import torch + strides = {'row_padding': (4, 1), 'column_gaps': (4, 2), + 'row_and_column_gaps': (7, 3)} + if (n_rows, num_experts, topk) != (3, 34, 2) or layout not in strides: + raise ValueError('Unknown strided bitmatrix control') + row_stride, col_stride = strides[layout] + storage_size = (n_rows - 1) * row_stride + (topk - 1) * col_stride + 1 + storage = torch.zeros(storage_size, dtype=torch.int16, device=device) + ids = storage.as_strided((n_rows, topk), (row_stride, col_stride)) + ids.copy_(torch.tensor([[0, 31], [1, 32], [2, 33]], + dtype=torch.int16, device=device)) + assert ids.stride() == (row_stride, col_stride) and not ids.is_contiguous() + return ids + + +EXTRA_CASES = [(17, 31, 1), (33, 32, 31), (513, 33, 32), + (5, 70, 33), (5, 70, 65), + (3, 34, 2, 'row_padding'), + (3, 34, 2, 'column_gaps'), + (3, 34, 2, 'row_and_column_gaps')] def run_control(index, load_module): if type(index) is not int or not 0 <= index < len(EXTRA_CASES): @@ -42,11 +82,14 @@ def run_control(index, load_module): except Exception as e: return (False, f'Failed to load module: {e}') device = 'cuda' - for i, (n_rows, num_experts, topk) in enumerate(EXTRA_CASES, start=0): + for i, control in enumerate(EXTRA_CASES, start=0): if i != index + 0: continue try: - topk_ids = make_boundary_topk_ids(n_rows, num_experts, topk, device) + n_rows, num_experts, topk = control[:3] + topk_ids = (make_strided_topk_ids(*control, device) + if len(control) == 4 else + make_boundary_topk_ids(n_rows, num_experts, topk, device)) result = mod.pack_topk_to_bitmatrix(topk_ids, num_experts) torch.cuda.synchronize() ref = reference_pack_bitmatrix(topk_ids.cpu(), num_experts).to(device) diff --git a/tasks/triton2triton/vllm/triton_pack_bitmatrix/config.yaml b/tasks/triton2triton/vllm/triton_pack_bitmatrix/config.yaml index f6a3ea271..f72956cb5 100644 --- a/tasks/triton2triton/vllm/triton_pack_bitmatrix/config.yaml +++ b/tasks/triton2triton/vllm/triton_pack_bitmatrix/config.yaml @@ -1,5 +1,7 @@ schema_version: 2 description: Optimize triton_pack_bitmatrix with the original Triton operator contract. +instructions: +- BOUNDARY_CHECKS.md candidate: language: triton initial_state: implemented diff --git a/tasks/triton2triton/vllm/triton_pack_bitmatrix/source/triton_pack_bitmatrix.py b/tasks/triton2triton/vllm/triton_pack_bitmatrix/source/triton_pack_bitmatrix.py index 6bd120c19..6049f5a6e 100644 --- a/tasks/triton2triton/vllm/triton_pack_bitmatrix/source/triton_pack_bitmatrix.py +++ b/tasks/triton2triton/vllm/triton_pack_bitmatrix/source/triton_pack_bitmatrix.py @@ -17,6 +17,7 @@ def pack_bitmatrix( n_expts_act, # num_topk BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, + N_CHUNKS: tl.constexpr, ): """ Pack topk_ids into a bitmatrix representation. @@ -32,15 +33,39 @@ def pack_bitmatrix( rem = indices % 32 one = tl.cast(1, tl.uint32) - for i in range(bm_cols): - offs = tl.arange(0, BLOCK_SIZE_K // 32) + i * (BLOCK_SIZE_K // 32) - x = tl.where( - mask[:, :, None] & (indices[:, :, None] >= 0) & (div[:, :, None] == offs[None, None, :]), - (one << rem)[:, :, None], 0 - ) - y = tl.reduce_or(x, axis=1) - bitmatrix_ptrs = bitmatrix + offsets_m[:, None] * bm_cols + offs[None, :] - tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows) + if N_CHUNKS == 1: + # Keep the original scored path intact when every assignment fits. + for i in range(bm_cols): + offs = tl.arange(0, BLOCK_SIZE_K // 32) + i * (BLOCK_SIZE_K // 32) + x = tl.where( + mask[:, :, None] & (indices[:, :, None] >= 0) & (div[:, :, None] == offs[None, None, :]), + (one << rem)[:, :, None], 0 + ) + y = tl.reduce_or(x, axis=1) + bitmatrix_ptrs = bitmatrix + offsets_m[:, None] * bm_cols + offs[None, :] + tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows) + else: + for i in range(bm_cols): + offs = tl.arange(0, 1) + i + x = tl.where( + mask[:, :, None] & (indices[:, :, None] >= 0) & (div[:, :, None] == offs[None, None, :]), + (one << rem)[:, :, None], 0 + ) + y = tl.reduce_or(x, axis=1) + # Fold every additional 32-assignment tile into this output word. + for chunk in range(1, N_CHUNKS): + chunk_k = chunk * BLOCK_SIZE_K + offsets_k + chunk_mask = (offsets_m < n_rows)[:, None] & (chunk_k < n_expts_act)[None, :] + chunk_offsets = offsets_m[:, None] * n_expts_act + chunk_k[None, :] + chunk_indices = tl.load(topk_ids + chunk_offsets, mask=chunk_mask, other=-1) + chunk_x = tl.where( + chunk_mask[:, :, None] & (chunk_indices[:, :, None] >= 0) + & (chunk_indices[:, :, None] // 32 == offs[None, None, :]), + (one << (chunk_indices % 32))[:, :, None], 0 + ) + y = y | tl.reduce_or(chunk_x, axis=1) + bitmatrix_ptrs = bitmatrix + offsets_m[:, None] * bm_cols + offs[None, :] + tl.store(bitmatrix_ptrs, y, mask=(offsets_m[:, None] < n_rows) & (offs[None, :] < bm_cols)) def pack_topk_to_bitmatrix( @@ -57,7 +82,9 @@ def pack_topk_to_bitmatrix( bitmatrix: [n_rows, bm_cols] uint32 tensor where each bit indicates if token is assigned to that expert """ - topk_ids = topk_ids.to(torch.int16) + # The kernel uses row-major linear offsets, so normalize strided views. + # This is a no-op for the contiguous inputs in the scored workloads. + topk_ids = topk_ids.to(torch.int16).contiguous() n_rows, num_topk = topk_ids.shape BLOCK_SIZE_M = 512 @@ -77,5 +104,6 @@ def pack_topk_to_bitmatrix( num_topk, BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_K=BLOCK_SIZE_K, + N_CHUNKS=triton.cdiv(num_topk, BLOCK_SIZE_K), ) return bitmatrix diff --git a/tasks/triton2triton/vllm/triton_pack_bitmatrix/workloads.json b/tasks/triton2triton/vllm/triton_pack_bitmatrix/workloads.json index fd7d6b31c..62d928d9a 100644 --- a/tasks/triton2triton/vllm/triton_pack_bitmatrix/workloads.json +++ b/tasks/triton2triton/vllm/triton_pack_bitmatrix/workloads.json @@ -174,6 +174,71 @@ "correctness" ], "status": "PASS" + }, + { + "test_case_id": "control-upstream-4", + "params": { + "case_index": 10003, + "configuration": [ + 5, + 70, + 33 + ], + "source": "repository" + }, + "checks": [ + "correctness" + ], + "status": "PASS" + }, + { + "test_case_id": "control-upstream-5", + "params": { + "case_index": 10004, + "configuration": [ + 5, + 70, + 65 + ], + "source": "repository" + }, + "checks": [ + "correctness" + ], + "status": "PASS" + }, + { + "test_case_id": "control-upstream-6", + "params": { + "case_index": 10005, + "configuration": [3, 34, 2], + "layout": "row_padding", + "source": "repository" + }, + "checks": ["correctness"], + "status": "PASS" + }, + { + "test_case_id": "control-upstream-7", + "params": { + "case_index": 10006, + "configuration": [3, 34, 2], + "layout": "column_gaps", + "source": "repository" + }, + "checks": ["correctness"], + "status": "PASS" + }, + { + "test_case_id": "control-upstream-8", + "params": { + "case_index": 10007, + "configuration": [3, 34, 2], + "layout": "row_and_column_gaps", + "source": "repository" + }, + "checks": ["correctness"], + "status": "PASS" } ], "migration": { @@ -195,6 +260,34 @@ 513, 33, 32 + ], + [ + 5, + 70, + 33 + ], + [ + 5, + 70, + 65 + ], + [ + 3, + 34, + 2, + "row_padding" + ], + [ + 3, + 34, + 2, + "column_gaps" + ], + [ + 3, + 34, + 2, + "row_and_column_gaps" ] ] } diff --git a/tasks/triton2triton/vllm/triton_temperature/README.md b/tasks/triton2triton/vllm/triton_temperature/README.md index 65454b994..d59a2942f 100644 --- a/tasks/triton2triton/vllm/triton_temperature/README.md +++ b/tasks/triton2triton/vllm/triton_temperature/README.md @@ -25,6 +25,13 @@ Canonical benchmark helpers must be materialized by Arena; do not edit their gen ## Protected evaluation controls -Temperature controls cover zero/one no-op, heterogeneous temperatures and remapped requests. - -The task-local `_arena_contract.py` and `_arena_replay.py` are protected evaluation code. Original cases, seeds, tolerances, warmups, sample counts, allocations and preparation boundaries remain in `scripts/task_runner.py`. The extra `contract_controls` manifest row is correctness-only. Both the frozen baseline and candidate receive the same checks. The measured graph exposes its real outputs; an untimed replay changes a domain-valid input, recomputes the CPU oracle and restores all input buffers in `finally`. For the zero operator the replay control instead poisons its output. References and snapshots are outside device timing. Failure to observe or replay the measured invocation is an error, never an accepted timing sample. +Temperature controls cover zero/one no-op, heterogeneous temperatures, remapped +requests at the largest scored vocabulary, a partial final 8192-element block +above the first block, and a padded row stride with untouched padding. The +partial-block control has nonidentity request mapping and rows whose final 257 +logits must change. The padded view is created on the target device so its +actual stride is checked at the candidate call. The supported layout is +contiguous columns with an optional padded row stride; arbitrary column strides +are outside this contract. + +The task-local `_arena_contract.py` and `_arena_replay.py` are protected evaluation code. Original cases, seeds, tolerances, warmups, sample counts, allocations and preparation boundaries remain in `scripts/task_runner.py`. The extra `contract_controls` manifest row is correctness-only. Both the frozen baseline and candidate receive the same checks. The measured graph exposes its real outputs; an untimed replay changes both temperature and request mapping, recomputes the CPU oracle and restores all input buffers in `finally`. References and snapshots are outside device timing. Failure to observe or replay the measured invocation is an error, never an accepted timing sample. diff --git a/tasks/triton2triton/vllm/triton_temperature/_arena_contract.py b/tasks/triton2triton/vllm/triton_temperature/_arena_contract.py index 0cb67985a..1436b26cc 100644 --- a/tasks/triton2triton/vllm/triton_temperature/_arena_contract.py +++ b/tasks/triton2triton/vllm/triton_temperature/_arena_contract.py @@ -16,6 +16,17 @@ def controls(harness, function, device): for args in control_inputs(harness): function(*to_device(args, device)) + # Build the view after allocating on the target device: transferring a CPU + # slice can make it contiguous and silently erase the row-stride control. + storage = torch.arange(4 * 264, dtype=torch.float32, device=device).reshape(4, 264) + 1 + logits = storage[:, :256] + mapping = i([2, 0, 3, 1]).to(device) + temperature = f([0, 1, 0.5, 2]).to(device) + assert logits.stride() == (264, 1) and logits.stride(0) > logits.shape[1] + padding = storage[:, 256:].clone() + function(logits, mapping, temperature) + torch.testing.assert_close(storage[:, 256:], padding, atol=0, rtol=0) + FUNCTION = 'apply_temperature' MUTABLE = (0,) ATOL = 0.01 @@ -33,10 +44,20 @@ def reference(harness, args): def fresh(args): args[2].add_(0.25) + args[1].copy_(args[1].roll(1)) def control_inputs(harness): yield (f([[1,-2,3,4], [2,1,-1,3], [4,6,-2,0], [2,3,4,5]]), i([2,0,3,1]), f([0,1,0.5,2])) + rows, vocab = 64, 32768 + yield (torch.linspace(-4, 4, rows * vocab, dtype=torch.float32).reshape(rows, vocab), + i([(row + 17) % rows for row in range(rows)]), + f([0, 1, 0.5, 2, 1.5, 0.75, 0.25, 3] * (rows // 8))) + # The large scored width is an exact multiple of the kernel's 8192-element + # block. This correctness-only width requires a masked second block. + rows, vocab = 4, 8192 + 257 + yield (torch.arange(1, rows * vocab + 1, dtype=torch.float32).reshape(rows, vocab), + i([2, 0, 3, 1]), f([0, 1, 0.5, 2])) def observe(result, args): return args[MUTABLE[0]] diff --git a/tasks/triton2triton/vllm/triton_temperature/_arena_eval.py b/tasks/triton2triton/vllm/triton_temperature/_arena_eval.py index ce87cd7c7..a84cfaaa2 100644 --- a/tasks/triton2triton/vllm/triton_temperature/_arena_eval.py +++ b/tasks/triton2triton/vllm/triton_temperature/_arena_eval.py @@ -3,6 +3,7 @@ import ast import copy import importlib.util +import inspect import json import math import os @@ -11,6 +12,7 @@ ROOT = Path(__file__).resolve().parent MANIFEST = ROOT / 'workloads.json' +PUBLIC_WRAPPER = 'apply_temperature(logits, idx_mapping, temperature)' def load_manifest(): @@ -28,6 +30,16 @@ def inspect_candidate(data, *, require_implemented=False): if not path.resolve().is_relative_to(ROOT): raise ValueError('Candidate path escapes task workspace') tree = ast.parse(path.read_text(), filename=source) + if source == 'source/triton_temperature.py': + wrappers = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == 'apply_temperature'] + if len(wrappers) != 1 or not isinstance(wrappers[0], ast.FunctionDef): + raise ValueError(f'Expected exactly one public {PUBLIC_WRAPPER} function') + args = wrappers[0].args + if (tuple(arg.arg for arg in args.args) != ('logits', 'idx_mapping', 'temperature') + or args.posonlyargs or args.kwonlyargs or args.vararg or args.kwarg + or args.defaults or args.kw_defaults): + raise ValueError(f'Public wrapper must have exact signature {PUBLIC_WRAPPER}') nodes = {n.name:n for n in tree.body if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))} for target in targets: node = nodes.get(target['name']) @@ -54,6 +66,30 @@ def load_harness(): spec = importlib.util.spec_from_file_location('_task_harness', path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) + original_load_module = module.load_module + + def load_module_with_public_signature_check(): + candidate = original_load_module() + exported = getattr(candidate, 'apply_temperature', None) + try: + signature = inspect.signature(exported, follow_wrapped=False) + except (TypeError, ValueError) as exc: + raise ValueError(f'Public wrapper must have exact signature {PUBLIC_WRAPPER}') from exc + params = tuple(signature.parameters.values()) + code = exported.__code__ if inspect.isfunction(exported) else None + if (not inspect.isfunction(exported) or len(params) != 3 + or tuple(param.name for param in params) != ('logits', 'idx_mapping', 'temperature') + or any(param.kind is not inspect.Parameter.POSITIONAL_OR_KEYWORD + or param.default is not inspect.Parameter.empty for param in params) + or code.co_posonlyargcount != 0 or code.co_argcount != 3 + or code.co_kwonlyargcount != 0 + or code.co_flags & (inspect.CO_VARARGS | inspect.CO_VARKEYWORDS) + or code.co_varnames[:3] != ('logits', 'idx_mapping', 'temperature') + or exported.__defaults__ or exported.__kwdefaults__): + raise ValueError(f'Public wrapper must have exact signature {PUBLIC_WRAPPER}') + return candidate + + module.load_module = load_module_with_public_signature_check import _arena_contract import _arena_replay _arena_replay.install(module, _arena_contract) @@ -76,6 +112,7 @@ def evaluate(role, action): for dependency in ('torch','triton'): if importlib.util.find_spec(dependency) is None: raise RuntimeError(f'Required runtime dependency unavailable: {dependency}') + harness.load_module() result['metadata']={'candidate_state':state,'input_table_verified':True} elif action == 'compile': ok, error = harness.run_compile() diff --git a/tasks/triton2triton/vllm/triton_temperature/workloads.json b/tasks/triton2triton/vllm/triton_temperature/workloads.json index 7757845ee..89166322a 100644 --- a/tasks/triton2triton/vllm/triton_temperature/workloads.json +++ b/tasks/triton2triton/vllm/triton_temperature/workloads.json @@ -125,7 +125,7 @@ ], "params": { "case_index": 5, - "coverage": "Temperature controls cover zero/one no-op, heterogeneous temperatures and remapped requests." + "coverage": "Correctness-only controls cover zero/one no-op, heterogeneous temperatures, remapped requests at the largest scored vocabulary, a partial final 8192-element block, and padded row strides." } } ], diff --git a/tests/test_combine_decode_controls.py b/tests/test_combine_decode_controls.py new file mode 100644 index 000000000..4816295d5 --- /dev/null +++ b/tests/test_combine_decode_controls.py @@ -0,0 +1,166 @@ +"""CPU checks for the combine task's correctness-only decode controls. + +These stand-ins exercise the task-owned correctness action; they do not qualify +the Triton kernel or its device timing. +""" + +import importlib.util +import os +from pathlib import Path +import sys +from types import SimpleNamespace + +import pytest +import torch + + +TASK = (Path(__file__).resolve().parents[1] / "tasks/triton2triton/vllm" + / "triton_combine_sampled_and_draft_tokens") + + +def load(path, name): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + old_cwd = os.getcwd() + try: + spec.loader.exec_module(module) + finally: + os.chdir(old_cwd) + return module + + +@pytest.fixture +def task_modules(monkeypatch): + monkeypatch.syspath_prepend(str(TASK)) + contract = load(TASK / "_arena_contract.py", "_combine_contract_test") + replay = load(TASK / "_arena_replay.py", "_combine_replay_test") + monkeypatch.setitem(sys.modules, "_arena_replay", replay) + harness = load(TASK / "scripts/task_runner.py", "_combine_harness_test") + adapter = load(TASK / "_arena_eval.py", "_combine_adapter_test") + return contract, replay, harness, adapter + + +def test_known_decode_answer_and_prefill_control(task_modules): + contract, _, harness, _ = task_modules + prefill, decode = list(contract.control_inputs(harness)) + + prefill_ids, prefill_indices = contract.reference(harness, prefill) + assert prefill_ids.tolist() == [-9, -9, 10, 1, 12, 7, 8] + assert prefill_indices.tolist() == [0, 2, 3, 4, 5, 6] + + assert decode[1].tolist() == [2, 0, 1] + assert decode[7].diff().tolist() == [1, 3, 2] + decode_ids, decode_indices = contract.reference(harness, decode) + assert decode_ids.tolist() == [303, 101, 111, 112, 202, 211] + assert decode_indices.tolist() == [0, 1, 2, 3, 4, 5] + + padded, storage = contract.padded_draft_inputs(harness, "cpu") + assert padded[6].stride() == (5, 1) + assert storage[:, 2:].tolist() == [[-777] * 3] * 3 + padded_ids, padded_indices = contract.reference(harness, padded) + assert torch.equal(padded_ids, decode_ids) + assert torch.equal(padded_indices, decode_indices) + + +@pytest.mark.parametrize("fault", [None, "omit_zero_draft_sampled", "use_batch_row", + "contiguous_draft_stride"]) +def test_native_correctness_action_rejects_decode_faults(task_modules, monkeypatch, fault): + contract, replay, harness, adapter = task_modules + observed_strides = [] + + def combine_sampled_and_draft_tokens( + input_ids, idx_mapping, last_sampled_tokens, query_start_loc, + seq_lens, prefill_len, draft_tokens, cu_num_logits, num_logits, + ): + observed_strides.append(draft_tokens.stride(0)) + original = input_ids.clone() + expected_ids, indices = harness.reference_combine( + input_ids, idx_mapping, last_sampled_tokens, query_start_loc, + seq_lens, prefill_len, draft_tokens, cu_num_logits, + ) + assert num_logits == indices.numel() + input_ids.copy_(expected_ids) + + if fault == "omit_zero_draft_sampled": + for batch in range(seq_lens.numel()): + state = idx_mapping[batch].item() + if (cu_num_logits[batch + 1] - cu_num_logits[batch] == 1 + and seq_lens[batch] > prefill_len[state]): + position = query_start_loc[batch + 1].item() - 1 + input_ids[position] = original[position] + elif fault == "use_batch_row": + for batch in range(seq_lens.numel()): + state = idx_mapping[batch].item() + if seq_lens[batch] <= prefill_len[state]: + continue + count = (cu_num_logits[batch + 1] - cu_num_logits[batch]).item() + start = query_start_loc[batch + 1].item() - count + input_ids[start] = last_sampled_tokens[batch] + if count > 1: + input_ids[start + 1:start + count] = draft_tokens[batch, :count - 1] + elif fault == "contiguous_draft_stride": + wrong_drafts = torch.as_strided( + draft_tokens, draft_tokens.shape, (draft_tokens.shape[1], 1)) + wrong_ids, _ = harness.reference_combine( + original, idx_mapping, last_sampled_tokens, query_start_loc, + seq_lens, prefill_len, wrong_drafts, cu_num_logits, + ) + input_ids.copy_(wrong_ids) + return indices + + if fault in ("omit_zero_draft_sampled", "use_batch_row"): + # Prove the new remapped decode input rejects each fault independently. + # Keep the normal control list intact for the native action below; + # padded_draft_inputs also reads its second entry. + decode = tuple(contract.control_inputs(harness))[1] + checked = replay.Recorder(harness, contract).wrap(combine_sampled_and_draft_tokens) + with pytest.raises(AssertionError, match="exact integer reference"): + checked(*decode) + + # Route the real task adapter through the normal installed correctness + # checker. The five scored cases are left to their separate GPU runner. + harness.load_module = lambda: SimpleNamespace( + combine_sampled_and_draft_tokens=combine_sampled_and_draft_tokens) + harness.run_correctness = lambda *, case_index=None: (True, None) + native_controls = contract.controls + monkeypatch.setattr(contract, "controls", + lambda harness, function, device: native_controls(harness, function, "cpu")) + replay.install(harness, contract) + monkeypatch.setattr(adapter, "load_harness", lambda: harness) + + result = adapter.evaluate("candidate", "correctness") + rows = {row["test_case_id"]: row for row in result["cases"]} + assert all(rows[f"perf{i}"]["status"] == "PASS" for i in range(1, 6)) + if fault is None: + assert result["status"] == rows["contract_controls"]["status"] == "PASS" + else: + assert result["status"] == rows["contract_controls"]["status"] == "FAIL" + assert "exact integer reference" in rows["contract_controls"]["reason"] + if fault == "contiguous_draft_stride": + assert observed_strides == [3, 2, 5] + + +@pytest.mark.skipif( + not torch.cuda.is_available() or importlib.util.find_spec("triton") is None, + reason="Triton and a compatible GPU are required", +) +def test_real_kernel_passes_native_correctness_action(task_modules, monkeypatch): + contract, replay, _, adapter = task_modules + monkeypatch.setitem(sys.modules, "_arena_contract", contract) + original_padded = contract.padded_draft_inputs + observed = [] + + def record_padded(harness, device): + args, storage = original_padded(harness, device) + observed.append((args[6].stride(), storage.device.type)) + return args, storage + + monkeypatch.setattr(contract, "padded_draft_inputs", record_padded) + + result = adapter.evaluate("candidate", "correctness") + assert result["status"] == "PASS", result + assert [row["test_case_id"] for row in result["cases"]] == [ + "perf1", "perf2", "perf3", "perf4", "perf5", "contract_controls", + ] + assert all(row["status"] == "PASS" for row in result["cases"]) + assert observed == [((5, 1), "cuda")] diff --git a/tests/test_pack_bitmatrix_benchmark_integrity.py b/tests/test_pack_bitmatrix_benchmark_integrity.py new file mode 100644 index 000000000..e1484d2bd --- /dev/null +++ b/tests/test_pack_bitmatrix_benchmark_integrity.py @@ -0,0 +1,282 @@ +"""The pack task must verify every invocation charged to a graph replay.""" + +import importlib.util +import os +from pathlib import Path +import shutil +import sys +from types import SimpleNamespace + +import pytest +import torch + +from src.perf_helper_materialization import materialize_perf_helpers_in_workspace + + +TASK = Path(__file__).resolve().parents[1] / "tasks/triton2triton/vllm/triton_pack_bitmatrix" + + +def _load(path, name): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + previous_directory = Path.cwd() + try: + spec.loader.exec_module(module) + finally: + os.chdir(previous_directory) + return module + + +class _TimedRun: + def _bind(self, replay, output): + self._replay = replay + self.outputs = output + + def rerun(self): + self.outputs = self._replay() + return self.outputs + + +def _fake_capture(mode="fresh", *, metadata_change=None, wrong_binding=False): + """Model four graph calls with real tensor writes and a bound replay.""" + controls = _load(TASK / "_upstream_controls.py", "_pack_integrity_controls") + checks = _load(TASK / "_arena_checks.py", "_pack_integrity_checks") + reference = controls.reference_pack_bitmatrix + state = {"calls": 0, "computes": 0, "capture_outputs": []} + harness = SimpleNamespace(reference_pack_bitmatrix=reference, _TimedRun=_TimedRun) + + def setup(): + topk_ids = torch.tensor( + [[0, 32], [1, 33]] if mode == "partial_overlap" else [[0, 33], [1, 34]], + dtype=torch.int16) + num_experts = 64 + shared = torch.zeros((2, 2), dtype=torch.uint32) + pool = torch.zeros((8, 2), dtype=torch.uint32) + interleaved_pool = torch.zeros((2, 8), dtype=torch.uint32) + overlapping_pool = torch.zeros((2, 5), dtype=torch.uint32) + + def candidate(ids, experts): + index = state["calls"] + state["calls"] += 1 + if mode == "alias": + if index == 3: + state["computes"] += 1 + shared.copy_(reference(ids, experts)) + return shared + state["computes"] += 1 + if mode == "pool": + output = pool[2 * index:2 * index + 2] + output.copy_(reference(ids, experts)) + return output + if mode == "interleaved_pool": + output = interleaved_pool[:, index::4] + output.copy_(reference(ids, experts)) + return output + if mode == "partial_overlap": + output = overlapping_pool[:, index:index + 2] + output.copy_(reference(ids, experts)) + return output + return reference(ids, experts) + + mod = SimpleNamespace(pack_topk_to_bitmatrix=candidate) + + def fn(): + return mod.pack_topk_to_bitmatrix(topk_ids, num_experts) + + def benchmark(measured, *, timed_run, **options): + del options + outputs = [measured() for _ in range(4)] + state["capture_outputs"] = outputs + + def replay(): + if mode == "alias": + shared.copy_(reference(topk_ids, num_experts)) + else: + targets = outputs[-1:] if mode == "stale_replay" else outputs + for output in targets: + output.copy_(reference(topk_ids, num_experts)) + return outputs[-1] + + timed_run._bind(replay, outputs[-1].clone() if wrong_binding else outputs[-1]) + metadata = {"benchmark_method": "cuda_graph", + "benchmark_timed_run_kind": "captured_graph", + "benchmark_effective_repeats": 4, + "benchmark_warmup": 10, + "benchmark_samples": 100} + if metadata_change: + metadata.update(metadata_change) + return 0.01, metadata + + return topk_ids, mod, candidate, fn, benchmark + + return checks, harness, state, setup() + + +@pytest.mark.parametrize("mode", ["fresh", "pool", "interleaved_pool"]) +def test_all_four_captured_outputs_and_perturbed_replay_pass(mode): + checks, harness, state, (ids, mod, candidate, fn, benchmark) = _fake_capture(mode) + pristine = ids.clone() + _, metadata = checks.checked_benchmark(harness, benchmark, fn) + assert metadata["timed_invocation_outputs_checked"] == 4 + assert metadata["timed_invocation_outputs_disjoint"] is True + assert metadata["timed_output_checked"] is True + assert metadata["perturbed_input_replay_checked"] is True + assert state["calls"] == state["computes"] == 4 + if mode in ("pool", "interleaved_pool"): + outputs = state["capture_outputs"] + assert len({output.untyped_storage().data_ptr() for output in outputs}) == 1 + if mode == "interleaved_pool": + spans = [checks.output_span(output) for output in outputs] + assert spans[0][0] < spans[1][1] and spans[1][0] < spans[0][1] + assert torch.equal(ids, pristine) + assert mod.pack_topk_to_bitmatrix is candidate + + +@pytest.mark.parametrize("mode,reason", [ + ("alias", "overlap"), + ("partial_overlap", "overlap"), + ("stale_replay", "independent reference"), +]) +def test_selected_compute_alias_and_distinct_stale_outputs_fail_and_restore(mode, reason): + checks, harness, state, (ids, mod, candidate, fn, benchmark) = _fake_capture(mode) + pristine = ids.clone() + with pytest.raises(AssertionError, match=reason): + checks.checked_benchmark(harness, benchmark, fn) + if mode == "alias": + assert state["calls"] == 4 and state["computes"] == 1 + assert len({id(output) for output in state["capture_outputs"]}) == 1 + assert torch.equal(ids, pristine) + assert mod.pack_topk_to_bitmatrix is candidate + + +@pytest.mark.parametrize("metadata_change,wrong_binding,reason", [ + ({"benchmark_effective_repeats": None}, False, "repeat count"), + ({"benchmark_effective_repeats": 5}, False, "not observable"), + ({"benchmark_timed_run_kind": "eager_callable"}, False, "not observable"), + (None, True, "final captured invocation"), +]) +def test_missing_or_mismatched_graph_binding_fails_closed_and_restores( + metadata_change, wrong_binding, reason, +): + checks, harness, _, (ids, mod, candidate, fn, benchmark) = _fake_capture( + metadata_change=metadata_change, wrong_binding=wrong_binding) + pristine = ids.clone() + with pytest.raises(AssertionError, match=reason): + checks.checked_benchmark(harness, benchmark, fn) + assert torch.equal(ids, pristine) + assert mod.pack_topk_to_bitmatrix is candidate + + +def test_event_fallback_keeps_one_observed_invocation(): + checks, harness, _, (ids, mod, candidate, fn, _) = _fake_capture() + pristine = ids.clone() + + def benchmark(measured, *, timed_run, **options): + del options + output = measured() + timed_run._bind(measured, output) + return 0.01, {"benchmark_method": "cuda_event_fallback", + "benchmark_timed_run_kind": "eager_callable", + "benchmark_effective_repeats": 1, + "benchmark_warmup": 10, + "benchmark_samples": 100} + + _, metadata = checks.checked_benchmark(harness, benchmark, fn) + assert metadata["timed_invocation_outputs_checked"] == 1 + assert metadata["perturbed_input_replay_checked"] is True + assert torch.equal(ids, pristine) + assert mod.pack_topk_to_bitmatrix is candidate + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton and a CUDA/ROCm GPU are required") +def test_actual_gpu_native_performance_and_output_ownership(tmp_path, request): + pytest.importorskip("triton") + previous_directory = Path.cwd() + previous_path = sys.path[:] + previous_helper = sys.modules.get("_aka_benchmark") + + def restore_import_state(): + os.chdir(previous_directory) + sys.path[:] = previous_path + if previous_helper is None: + sys.modules.pop("_aka_benchmark", None) + else: + sys.modules["_aka_benchmark"] = previous_helper + + request.addfinalizer(restore_import_state) + task = tmp_path / "pack_task" + shutil.copytree(TASK, task) + materialize_perf_helpers_in_workspace(task) + adapter = _load(task / "_arena_eval.py", "_pack_integrity_gpu_adapter") + result = adapter.evaluate("candidate", "performance") + assert result["status"] == "PASS", result + assert [row["test_case_id"] for row in result["cases"]] == [f"perf{i}" for i in range(1, 6)] + for row in result["cases"]: + assert row["status"] == "PASS", row + measurement = row["metadata"]["harness_measurement"] + assert measurement["benchmark_method"] == "cuda_graph" + assert measurement["benchmark_warmup"] == 10 + assert measurement["benchmark_samples"] == 100 + assert measurement["timed_invocation_outputs_checked"] == measurement["benchmark_effective_repeats"] + assert measurement["timed_invocation_outputs_disjoint"] is True + + checks = _load(task / "_arena_checks.py", "_pack_integrity_gpu_checks") + harness = adapter.load_harness() + source = _load(task / "source/triton_pack_bitmatrix.py", "_pack_integrity_gpu_source") + pool_capacity = 32 + pool = torch.empty((32, pool_capacity * 2), dtype=torch.uint32, device="cuda") + pool_calls = [0] + + def interleaved(ids, experts): + index = pool_calls[0] + pool_calls[0] += 1 + assert index < pool_capacity + output = pool[:, index::pool_capacity] + output.copy_(source.pack_topk_to_bitmatrix(ids, experts)) + return output + + def setup_interleaved(): + topk_ids = torch.randint(0, 64, (32, 2), device="cuda", dtype=torch.int16) + num_experts = 64 + mod = SimpleNamespace(pack_topk_to_bitmatrix=interleaved) + + def fn(): + mod.pack_topk_to_bitmatrix(topk_ids, num_experts) + + return fn, topk_ids + + interleaved_fn, interleaved_ids = setup_interleaved() + pristine_ids = interleaved_ids.clone() + _, interleaved_metadata = checks.checked_benchmark( + harness, harness._benchmark_cuda_graph_or_events, interleaved_fn, + warmup=1, repetition=3, estimate_reps=1, target_ms=0.2, max_graph_repeats=8, + ) + assert interleaved_metadata["benchmark_method"] == "cuda_graph" + assert interleaved_metadata["benchmark_effective_repeats"] >= 2 + assert (interleaved_metadata["timed_invocation_outputs_checked"] + == interleaved_metadata["benchmark_effective_repeats"]) + assert interleaved_metadata["timed_invocation_outputs_disjoint"] is True + assert interleaved_metadata["perturbed_input_replay_checked"] is True + assert torch.equal(interleaved_ids, pristine_ids) + + shared = torch.zeros((32, 2), dtype=torch.uint32, device="cuda") + + def aliased(ids, experts): + shared.copy_(source.pack_topk_to_bitmatrix(ids, experts)) + return shared + + def setup(): + topk_ids = torch.randint(0, 64, (32, 2), device="cuda", dtype=torch.int16) + num_experts = 64 + mod = SimpleNamespace(pack_topk_to_bitmatrix=aliased) + + def fn(): + mod.pack_topk_to_bitmatrix(topk_ids, num_experts) + + return fn + + with pytest.raises(AssertionError, match="overlap"): + checks.checked_benchmark( + harness, harness._benchmark_cuda_graph_or_events, setup(), + warmup=1, repetition=3, target_ms=0.2, max_graph_repeats=8, + ) diff --git a/tests/test_pack_bitmatrix_topk_boundary.py b/tests/test_pack_bitmatrix_topk_boundary.py new file mode 100644 index 000000000..0f82d6c78 --- /dev/null +++ b/tests/test_pack_bitmatrix_topk_boundary.py @@ -0,0 +1,223 @@ +"""Native and device regressions for assignments beyond the first 32 top-k slots.""" + +import importlib.util +import json +import os +from pathlib import Path +import sys +from types import SimpleNamespace + +import pytest +import torch + +from src.task_prompt import build_task_prompt +from src.task_spec import load_task_spec + + +TASK = Path(__file__).resolve().parents[1] / "tasks/triton2triton/vllm/triton_pack_bitmatrix" + + +def _load(path, name): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + previous_directory = Path.cwd() + try: + spec.loader.exec_module(module) + finally: + os.chdir(previous_directory) + return module + + +def _assert_tile_signatures(controls, assignments, experts): + """Every extra tile must add a bit, and different rows must stay different.""" + reference = controls.reference_pack_bitmatrix + full = reference(assignments, experts) + first = reference(assignments[:, :32], experts) + assert full.shape == (5, (experts + 31) // 32) + assert all(not torch.equal(full[row], first[row]) for row in range(5)) + assert len({tuple(row) for row in full.tolist()}) == 5 + assert assignments[:, -1].tolist() == [experts - 1 - row for row in range(5)] + assert torch.equal(assignments[:, 0], assignments[:, 1]) + if assignments.shape[1] == 65: + first_two = reference(assignments[:, :64], experts) + missing_middle = reference(torch.cat((assignments[:, :32], assignments[:, 64:]), dim=1), experts) + assert assignments[:, 32].tolist() == [33, 34, 35, 36, 37] + assert all(not torch.equal(first_two[row], first[row]) for row in range(5)) + assert all(not torch.equal(full[row], first_two[row]) for row in range(5)) + assert all(not torch.equal(full[row], missing_middle[row]) for row in range(5)) + return full + + +@pytest.mark.parametrize("case_index,topk", [(10003, 33), (10004, 65)]) +def test_native_controls_reject_tile_and_row_mutants(monkeypatch, case_index, topk): + controls = _load(TASK / "_upstream_controls.py", "_pack_topk_controls") + runner = _load(TASK / "scripts/task_runner.py", "_pack_topk_runner") + adapter = _load(TASK / "_arena_eval.py", "_pack_topk_adapter") + monkeypatch.setitem(sys.modules, "_upstream_controls", controls) + original_generator = controls.make_boundary_topk_ids + monkeypatch.setattr( + controls, "make_boundary_topk_ids", + lambda rows, experts, count, device: original_generator(rows, experts, count, "cpu"), + ) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + original_to = torch.Tensor.to + monkeypatch.setattr( + torch.Tensor, "to", + lambda tensor, *args, **kwargs: tensor if args == ("cuda",) + else original_to(tensor, *args, **kwargs), + ) + + manifest = adapter.load_manifest() + rows = {row["params"]["case_index"]: row for row in manifest["cases"] + if row["test_case_id"].startswith("control-upstream-")} + assert rows[case_index]["params"]["configuration"] == [5, 70, topk] + assert rows[case_index]["params"]["source"] == "repository" + assert rows[case_index]["checks"] == ["correctness"] + assert all("source_revision" in rows[10000 + i]["params"] for i in range(3)) + assert manifest["upstream_controls"] == json.loads(json.dumps(controls.EXTRA_CASES)) + assert runner.TEST_SHAPES == [(32, 8, 2), (64, 16, 2), (128, 32, 2), + (256, 64, 4), (512, 64, 2)] + assert (runner.WARMUP_ITERATIONS, runner.BENCHMARK_ITERATIONS) == (10, 100) + + assignments = original_generator(5, 70, topk, "cpu") + assert assignments.dtype == torch.int16 + assert {0, 31, 32, 63, 64, 69} <= set(assignments.flatten().tolist()) + expected = _assert_tile_signatures(controls, assignments, 70) + assert torch.all(expected[:, 2] != 0) + + reference = controls.reference_pack_bitmatrix + good = SimpleNamespace(pack_topk_to_bitmatrix=reference) + monkeypatch.setattr(runner, "load_module", lambda: good) + assert runner.run_correctness(case_index=case_index) == (True, None) + + def wrong_extra_row(ids, experts): + wrong = ids.clone() + wrong[1:, 32:] = ids[0, 32:] + return reference(wrong, experts) + + mutants = { + "first-32-only": lambda ids, experts: reference(ids[:, :32], experts), + "wrong-extra-row": wrong_extra_row, + } + if topk == 65: + mutants["missing-middle"] = lambda ids, experts: reference( + torch.cat((ids[:, :32], ids[:, 64:]), dim=1), experts) + for name, mutant in mutants.items(): + monkeypatch.setattr(runner, "load_module", lambda mutant=mutant: SimpleNamespace( + pack_topk_to_bitmatrix=mutant)) + passed, reason = runner.run_correctness(case_index=case_index) + assert not passed and "mismatched elements" in reason, (name, reason) + + +@pytest.mark.parametrize("topk,experts", [(33, 64), (65, 64), (33, 70), (65, 70)]) +def test_kernel_packs_distinct_tiles_rows_and_output_word_tail_on_gpu(monkeypatch, topk, experts): + if not torch.cuda.is_available(): + pytest.skip("A CUDA or ROCm GPU is required") + pytest.importorskip("triton") + controls = _load(TASK / "_upstream_controls.py", "_pack_topk_gpu_controls") + kernel = _load(TASK / "source/triton_pack_bitmatrix.py", "_pack_topk_gpu_kernel") + assignments = controls.make_boundary_topk_ids(5, experts, topk, "cpu") + expected = _assert_tile_signatures(controls, assignments, experts) + + actual = kernel.pack_topk_to_bitmatrix(assignments.cuda(), experts) + assert actual.shape == (5, (experts + 31) // 32) + assert actual.dtype == torch.uint32 + assert torch.equal(actual, expected.cuda()) + + if experts == 70: + runner = _load(TASK / "scripts/task_runner.py", "_pack_topk_gpu_runner") + monkeypatch.setitem(sys.modules, "_upstream_controls", controls) + passed, reason = runner.run_correctness(case_index=10003 if topk == 33 else 10004) + assert passed, reason + + +@pytest.mark.parametrize("case_index,layout,strides", [ + (10005, "row_padding", (4, 1)), + (10006, "column_gaps", (4, 2)), + (10007, "row_and_column_gaps", (7, 3)), +]) +def test_native_strided_controls_reject_contiguous_addressing( + monkeypatch, case_index, layout, strides, +): + controls = _load(TASK / "_upstream_controls.py", "_pack_stride_controls") + runner = _load(TASK / "scripts/task_runner.py", "_pack_stride_runner") + adapter = _load(TASK / "_arena_eval.py", "_pack_stride_adapter") + monkeypatch.setitem(sys.modules, "_upstream_controls", controls) + original_generator = controls.make_strided_topk_ids + monkeypatch.setattr(controls, "make_strided_topk_ids", + lambda rows, experts, count, kind, device: + original_generator(rows, experts, count, kind, "cpu")) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + original_to = torch.Tensor.to + monkeypatch.setattr(torch.Tensor, "to", + lambda tensor, *args, **kwargs: tensor if args == ("cuda",) + else original_to(tensor, *args, **kwargs)) + + manifest = adapter.load_manifest() + row = next(r for r in manifest["cases"] if r["params"].get("case_index") == case_index) + assert row["params"]["configuration"] == [3, 34, 2] + assert row["params"]["layout"] == layout + assert row["params"]["source"] == "repository" + assert row["checks"] == ["correctness"] + assert len(manifest["cases"]) == 13 + ids = original_generator(3, 34, 2, layout, "cpu") + assert ids.stride() == strides and not ids.is_contiguous() + assert ids.tolist() == [[0, 31], [1, 32], [2, 33]] + reference = controls.reference_pack_bitmatrix + expected = reference(ids, 34) + assert expected.tolist() == [[2147483649, 0], [2, 1], [4, 2]] + + good = SimpleNamespace(pack_topk_to_bitmatrix=reference) + monkeypatch.setattr(runner, "load_module", lambda: good) + assert runner.run_correctness(case_index=case_index) == (True, None) + + def wrong_linear_addressing(source, experts): + wrong_view = source.as_strided(source.shape, (source.shape[1], 1)) + return reference(wrong_view, experts) + + assert not torch.equal(wrong_linear_addressing(ids, 34), expected) + bad = SimpleNamespace(pack_topk_to_bitmatrix=wrong_linear_addressing) + monkeypatch.setattr(runner, "load_module", lambda: bad) + passed, reason = runner.run_correctness(case_index=case_index) + assert not passed and "mismatched elements" in reason + + +@pytest.mark.parametrize("case_index,layout,strides", [ + (10005, "row_padding", (4, 1)), + (10006, "column_gaps", (4, 2)), + (10007, "row_and_column_gaps", (7, 3)), +]) +def test_gpu_strided_inputs_reject_old_linear_load_and_pass_wrapper( + monkeypatch, case_index, layout, strides, +): + if not torch.cuda.is_available(): + pytest.skip("A CUDA or ROCm GPU is required") + triton = pytest.importorskip("triton") + controls = _load(TASK / "_upstream_controls.py", "_pack_stride_gpu_controls") + source = _load(TASK / "source/triton_pack_bitmatrix.py", "_pack_stride_gpu_source") + ids = controls.make_strided_topk_ids(3, 34, 2, layout, "cuda") + assert ids.stride() == strides and not ids.is_contiguous() + pristine = ids.clone() + expected = controls.reference_pack_bitmatrix(ids.cpu(), 34).cuda() + + # Calling the unchanged kernel directly with a strided input reproduces + # its former linear-addressing bug; the public wrapper must normalize it. + old_output = torch.zeros((3, 2), device="cuda", dtype=torch.uint32) + source.pack_bitmatrix[(triton.cdiv(3, 512),)]( + old_output, ids, 3, 2, 2, BLOCK_SIZE_M=512, BLOCK_SIZE_K=32, N_CHUNKS=1) + assert not torch.equal(old_output, expected) + actual = source.pack_topk_to_bitmatrix(ids, 34) + assert torch.equal(actual, expected) + assert torch.equal(ids, pristine) + + runner = _load(TASK / "scripts/task_runner.py", "_pack_stride_gpu_runner") + monkeypatch.setitem(sys.modules, "_upstream_controls", controls) + assert runner.run_correctness(case_index=case_index) == (True, None) + + +def test_boundary_instructions_are_in_actual_task_prompt(): + spec = load_task_spec(TASK / "config.yaml", task_id="triton2triton/vllm/triton_pack_bitmatrix") + assert "BOUNDARY_CHECKS.md" in spec.to_mapping()["instructions"] + instructions = (TASK / "BOUNDARY_CHECKS.md").read_text() + prompt = build_task_prompt(spec, TASK, target_gpu="MI355X") + assert f"Task instructions from BOUNDARY_CHECKS.md:\n{instructions}" in prompt diff --git a/tests/test_temperature_layout_controls.py b/tests/test_temperature_layout_controls.py new file mode 100644 index 000000000..de7455eb1 --- /dev/null +++ b/tests/test_temperature_layout_controls.py @@ -0,0 +1,309 @@ +"""Temperature task controls; CPU checks do not qualify the GPU task.""" + +import importlib.util +import json +import os +from pathlib import Path +import sys + +import pytest +import torch + + +TASK = Path(__file__).resolve().parents[1] / "tasks/triton2triton/vllm/triton_temperature" + + +def _load(path, name): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + cwd = os.getcwd() + try: + spec.loader.exec_module(module) + finally: + os.chdir(cwd) + return module + + +@pytest.fixture +def task_modules(monkeypatch): + replay = _load(TASK / "_arena_replay.py", "temperature_layout_replay") + contract = _load(TASK / "_arena_contract.py", "temperature_layout_contract") + monkeypatch.setitem(sys.modules, "_arena_replay", replay) + return replay, contract + + +def _correct(contract, logits, mapping, temperature): + logits.copy_(contract.reference(None, (logits, mapping, temperature))) + + +def test_manifest_keeps_five_original_scored_cases(): + manifest = json.loads((TASK / "workloads.json").read_text()) + assert manifest["input_table"] == [[4, 256], [8, 1024], [16, 4096], + [32, 8192], [64, 32768]] + assert [(row["test_case_id"], row["checks"]) for row in manifest["cases"]] == [ + (f"perf{index}", ["correctness", "performance"]) for index in range(1, 6) + ] + [("contract_controls", ["correctness"])] + assert "partial final 8192-element block" in manifest["cases"][-1]["params"]["coverage"] + + +@pytest.mark.parametrize("replacement", [ + "def apply_temperature(logits, idx_mapping, temperature, extra=1):", + "def apply_temperature(logits, idx_mapping, temperature, *args):", + "def apply_temperature(logits, idx_mapping, temperature, **kwargs):", + "def apply_temperature(logits, idx_mapping, temperature=1):", + "def apply_temperature(logits, idx_mapping, *, temperature):", + "def apply_temperature(logits, idx_mapping, /, temperature):", + "def apply_temperature(logits, temperature, idx_mapping):", + "def apply_temperature(logits, mapping, temperature):", +]) +def test_public_wrapper_signature_mutations_fail_before_harness(replacement, tmp_path, monkeypatch): + adapter = _load(TASK / "_arena_eval.py", "temperature_signature_adapter") + original = "def apply_temperature(logits, idx_mapping, temperature):" + source = (TASK / "source/triton_temperature.py").read_text() + assert source.count(original) == 1 + target = tmp_path / "source/triton_temperature.py" + target.parent.mkdir() + target.write_text(source.replace(original, replacement, 1)) + monkeypatch.setattr(adapter, "ROOT", tmp_path) + + def unexpected_harness(): + pytest.fail("invalid signature reached the GPU harness") + + monkeypatch.setattr(adapter, "load_harness", unexpected_harness) + result = adapter.evaluate("candidate", "compile") + assert result["status"] == "FAIL" + assert result["failure_kind"] == "execution_failure" + assert "exact signature apply_temperature(logits, idx_mapping, temperature)" in result["reason"] + + +def test_public_wrapper_signature_accepts_original_and_ignores_annotations(tmp_path, monkeypatch): + adapter = _load(TASK / "_arena_eval.py", "temperature_signature_adapter_valid") + data = adapter.load_manifest() + assert adapter.inspect_candidate(data, require_implemented=True) == "implemented" + original = "def apply_temperature(logits, idx_mapping, temperature):" + annotated = "def apply_temperature(logits: object, idx_mapping: object, temperature: object) -> object:" + source = (TASK / "source/triton_temperature.py").read_text() + target = tmp_path / "source/triton_temperature.py" + target.parent.mkdir() + target.write_text(source.replace(original, annotated, 1)) + monkeypatch.setattr(adapter, "ROOT", tmp_path) + assert adapter.inspect_candidate(data, require_implemented=True) == "implemented" + + +@pytest.mark.parametrize("wrapper_source, accepted", [ + ("def apply_temperature(logits, idx_mapping, temperature):\n return None\n", True), + ("def apply_temperature(logits, idx_mapping, temperature):\n return None\n" + "apply_temperature.__defaults__ = ()\n" + "apply_temperature.__kwdefaults__ = {}\n", True), + ("def apply_temperature(logits, idx_mapping, temperature):\n return None\n" + "_original_apply_temperature = apply_temperature\n" + "apply_temperature = lambda logits, idx_mapping, temperature, extra=None: " + "_original_apply_temperature(logits, idx_mapping, temperature)\n", False), + ("import inspect\n" + "def apply_temperature(logits, idx_mapping, temperature):\n return None\n" + "_original_apply_temperature = apply_temperature\n" + "apply_temperature = lambda logits, idx_mapping, temperature, extra=None: " + "_original_apply_temperature(logits, idx_mapping, temperature)\n" + "apply_temperature.__signature__ = inspect.signature(_original_apply_temperature)\n", False), + ("import functools\n" + "def variadic(fn):\n" + " @functools.wraps(fn)\n" + " def wrapper(*args, **kwargs):\n" + " return fn(*args, **kwargs)\n" + " return wrapper\n" + "@variadic\n" + "def apply_temperature(logits, idx_mapping, temperature):\n" + " return None\n", False), +]) +def test_exported_wrapper_signature_checked_by_native_loader( + wrapper_source, accepted, task_modules, tmp_path, monkeypatch): + replay, contract = task_modules + monkeypatch.setitem(sys.modules, "_arena_contract", contract) + adapter = _load(TASK / "_arena_eval.py", "temperature_exported_signature_adapter") + source_path = tmp_path / "source/triton_temperature.py" + source_path.parent.mkdir() + source_path.write_text( + "class triton:\n" + " @staticmethod\n" + " def jit(fn): return fn\n" + "@triton.jit\n" + "def _temperature_kernel(): return None\n" + + wrapper_source) + runner_path = tmp_path / "scripts/task_runner.py" + runner_path.parent.mkdir() + runner_path.write_bytes((TASK / "scripts/task_runner.py").read_bytes()) + monkeypatch.setattr(adapter, "ROOT", tmp_path) + assert adapter.inspect_candidate(adapter.load_manifest(), require_implemented=True) == "implemented" + cwd = os.getcwd() + try: + result = adapter.evaluate("candidate", "compile") + finally: + os.chdir(cwd) + assert result["status"] == ("PASS" if accepted else "FAIL") + if not accepted: + assert result["failure_kind"] == "execution_failure" + assert "exact signature apply_temperature(logits, idx_mapping, temperature)" in result["reason"] + + +def test_native_controls_require_large_mapping_partial_block_and_padded_stride(task_modules): + replay, contract = task_modules + seen = [] + padded = [] + + def correct(logits, mapping, temperature): + seen.append((tuple(logits.shape), logits.stride(), mapping.tolist())) + before = logits.clone() if logits.stride(0) > logits.shape[1] else None + _correct(contract, logits, mapping, temperature) + if before is not None: + padded.append((before, logits.clone())) + + contract.controls(None, replay.Recorder(None, contract).wrap(correct), "cpu") + assert seen[1][0] == (64, 32768) + assert seen[1][2] != list(range(64)) + assert seen[2] == ((4, 8449), (8449, 1), [2, 0, 3, 1]) + assert seen[3][0] == (4, 256) + assert seen[3][1] == (264, 1) + tail = list(contract.control_inputs(None))[2] + expected = contract.reference(None, tail) + torch.testing.assert_close(expected[0, 8192:], tail[0][0, 8192:] * 2, atol=0, rtol=0) + torch.testing.assert_close(expected[2, 8192:], tail[0][2, 8192:] / 2, atol=0, rtol=0) + assert not torch.equal(expected[0, 8192:], tail[0][0, 8192:]) + assert not torch.equal(expected[2, 8192:], tail[0][2, 8192:]) + before, after = padded[0] + torch.testing.assert_close(after[0], before[0] * 2, atol=0, rtol=0) + torch.testing.assert_close(after[1], before[1], atol=0, rtol=0) + torch.testing.assert_close(after[2], before[2] / 2, atol=0, rtol=0) + torch.testing.assert_close(after[3], before[3], atol=0, rtol=0) + + +def test_native_controls_reject_candidate_skipping_partial_final_block(task_modules): + replay, contract = task_modules + seen = [] + + def skip_partial_final_block(logits, mapping, temperature): + seen.append(tuple(logits.shape)) + if logits.shape[1] > 8192 and logits.shape[1] % 8192: + _correct(contract, logits[:, :8192], mapping, temperature) + else: + _correct(contract, logits, mapping, temperature) + + # The mutation satisfies the existing small and 32768-wide controls. + for args in list(contract.control_inputs(None))[:2]: + replay.Recorder(None, contract).wrap(skip_partial_final_block)(*args) + assert seen == [(4, 4), (64, 32768)] + + with pytest.raises(AssertionError): + contract.controls(None, replay.Recorder(None, contract).wrap(skip_partial_final_block), "cpu") + assert seen[-1] == (4, 8449) + + +def test_ignored_mapping_fails_large_vocab_control(task_modules): + replay, contract = task_modules + large = list(contract.control_inputs(None))[1] + + def ignore_mapping(logits, mapping, temperature): + identity = torch.arange(mapping.numel(), dtype=mapping.dtype) + _correct(contract, logits, identity, temperature) + + wrapped = replay.Recorder(None, contract).wrap(ignore_mapping) + with pytest.raises(AssertionError): + wrapped(*large) + + +def test_contiguous_row_assumption_and_padding_corruption_fail(task_modules): + replay, contract = task_modules + + def assume_contiguous(logits, mapping, temperature): + if logits.stride(0) > logits.shape[1]: + flat_rows = torch.as_strided(logits, logits.shape, (logits.shape[1], 1)) + _correct(contract, flat_rows, mapping, temperature) + else: + _correct(contract, logits, mapping, temperature) + + with pytest.raises(AssertionError): + contract.controls(None, replay.Recorder(None, contract).wrap(assume_contiguous), "cpu") + + def corrupt_padding(logits, mapping, temperature): + _correct(contract, logits, mapping, temperature) + if logits.stride(0) > logits.shape[1]: + all_columns = torch.as_strided(logits, (logits.shape[0], logits.stride(0)), + logits.stride()) + all_columns[:, logits.shape[1]:].fill_(-999) + + with pytest.raises(AssertionError): + contract.controls(None, replay.Recorder(None, contract).wrap(corrupt_padding), "cpu") + + +def test_replay_varies_mapping_and_restores_it(task_modules): + replay, contract = task_modules + args = next(contract.control_inputs(None)) + before = replay.clone(args) + contract.fresh(args) + assert not torch.equal(args[1], before[1]) + assert not torch.equal(args[2], before[2]) + replay.restore(args, before) + for actual, expected in zip(args, before): + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA/ROCm GPU") +def test_actual_gpu_kernel_passes_native_controls_and_mapping_replay(task_modules, monkeypatch): + pytest.importorskip("triton") + replay, contract = task_modules + harness = _load(TASK / "scripts/task_runner.py", "temperature_layout_harness") + original_loader = harness.load_module + seen = [] + + def load_with_stride_probe(): + module = original_loader() + actual = module.apply_temperature + + def inspected(logits, idx_mapping, temperature): + seen.append((tuple(logits.shape), logits.stride())) + return actual(logits, idx_mapping, temperature) + + module.apply_temperature = inspected + return module + + harness.load_module = load_with_stride_probe + replay.install(harness, contract) + ok, error = harness.run_correctness(case_index=contract.CONTROL_INDEX) + assert ok, error + assert ((64, 32768), (32768, 1)) in seen + assert ((4, 8449), (8449, 1)) in seen + assert ((4, 256), (264, 1)) in seen + + source = original_loader() + logits = torch.arange(8 * 1024, dtype=torch.float32, device="cuda").reshape(8, 1024) + 1 + mapping = torch.arange(8, dtype=torch.int32, device="cuda") + temperature = torch.tensor([0.5, 2, 1.5, 0.75, 3, 0.25, 1.25, 2.5], device="cuda") + args = (logits, mapping, temperature) + pristine = replay.clone(args) + replay_mappings = [] + + def observed(logits, idx_mapping, temperature): + replay_mappings.append(idx_mapping.cpu().clone()) + return source.apply_temperature(logits, idx_mapping, temperature) + + class SingleRun: + def _bind(self, run, output): + self.run, self.outputs = run, output + + def rerun(self): + self.outputs = self.run() + return self.outputs + + def one_gpu_call(run, *, timed_run, **_options): + timed_run._bind(run, run()) + return 1.0, {"benchmark_method": "single_gpu_call_test"} + + # The committed task runner has only the generated benchmark stub. This + # single-invocation regression supplies its own replay holder. + monkeypatch.setattr(harness, "_TimedRun", SingleRun, raising=False) + recorder = replay.Recorder(harness, contract) + function = recorder.wrap(observed) + _, metadata = recorder.benchmark(one_gpu_call, lambda: function(*args)) + assert metadata["timed_output_checked"] and metadata["replay_input_control_checked"] + assert any(not torch.equal(item, pristine[1].cpu()) for item in replay_mappings) + for actual, expected in zip(args, pristine): + assert torch.equal(actual, expected) diff --git a/tests/test_triton_task_migration_v2.py b/tests/test_triton_task_migration_v2.py index f21e940af..938ddbfdd 100644 --- a/tests/test_triton_task_migration_v2.py +++ b/tests/test_triton_task_migration_v2.py @@ -189,11 +189,34 @@ def test_vllm_v2_preserves_all_original_cases_checks_sources_and_helpers(path): f'{operand.upper()}_LARGE={operand}.numel() > 2**31'.encode(), f'{operand.upper()}_LARGE=_requires_int64_index({operand})'.encode()) if task.name == 'triton_pack_bitmatrix': - old = b'div[:, :, None] == offs[None, None, :], (one << rem)[:, :, None], 0' - new = (b'mask[:, :, None] & (indices[:, :, None] >= 0) & (div[:, :, None] == offs[None, None, :]),\n' - b' (one << rem)[:, :, None], 0') - assert original.count(old) == 1 - original = original.replace(old, new) + # Explicit semantic repairs: ignore tile padding and OR every + # 32-assignment tile, including top-k 33/65 tails. The dedicated + # boundary tests reject truncation. The wrapper adds the tile count + # and normalizes strided inputs; retain every other source byte. + current = source.read_text() + before_tree, after_tree = ast.parse(original), ast.parse(current) + before_kernel = next(n for n in before_tree.body + if isinstance(n, ast.FunctionDef) and n.name == 'pack_bitmatrix') + after_kernel = next(n for n in after_tree.body + if isinstance(n, ast.FunctionDef) and n.name == 'pack_bitmatrix') + assert hashlib.sha256(ast.dump(after_kernel, include_attributes=False).encode()).hexdigest() == ( + 'bcb27618446d9a98ce504268fbd520b56802d7baa44ac7a4f9cdaebc12d18cad' + ) + original = original.replace(ast.get_source_segment(original.decode(), before_kernel).encode(), + ast.get_source_segment(current, after_kernel).encode(), 1) + old_launch = b' BLOCK_SIZE_K=BLOCK_SIZE_K,\n )' + new_launch = (b' BLOCK_SIZE_K=BLOCK_SIZE_K,\n' + b' N_CHUNKS=triton.cdiv(num_topk, BLOCK_SIZE_K),\n )') + assert original.count(old_launch) == 1 + original = original.replace(old_launch, new_launch) + old_layout = b' topk_ids = topk_ids.to(torch.int16)' + new_layout = ( + b' # The kernel uses row-major linear offsets, so normalize strided views.\n' + b' # This is a no-op for the contiguous inputs in the scored workloads.\n' + b' topk_ids = topk_ids.to(torch.int16).contiguous()' + ) + assert original.count(old_layout) == 1 + original = original.replace(old_layout, new_layout) if task.name == 'triton_batched_moe': # Only add the missing N-tail load mask; preserve all other source # bytes. Dedicated MoE controls cover inactive experts and tails. @@ -3986,7 +4009,9 @@ def replay(): if mode == 'mutate_replay': ids.zero_() return output timed_run.outputs, timed_run.rerun = output, replay - return .125, {'benchmark_method': 'cuda_graph'} + return .125, {'benchmark_method': 'cuda_graph', + 'benchmark_timed_run_kind': 'captured_graph', + 'benchmark_effective_repeats': 1} h._benchmark_cuda_graph_or_events = benchmark checks.install(h) rows = h.run_performance()