Fix boundary coverage and top-k packing in three Triton tasks

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.

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

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."
       }
     }
   ],

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.

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)

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)

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

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