[Perf] Reuse qualified QSA and PLE prefill stages on 9683

Port the peer TP4 QSA row partition and bounded PLE prefill convolution, retaining tile lifetime repairs, explicit admission and ordinary decode paths. Record destination operator gates and the complete 15.217 ms verifier window.

Assisted-by: Codex
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>

benchmarks/flashnext_prefill/ple_prefill_conv.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Screen PLE prefill history elimination, including state commit and SiLU.
+
+Before candidate execution: finite outputs and maximum row relative-L2 <=
+0.001 versus the production convolution; state must match exactly, including
+unused/speculative slots. This is a microbenchmark, not model qualification.
+"""
+
+import argparse
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+from moe_tiles import measure, relative_l2
+
+from vllm.model_executor.kernels.ple.short_conv import dilated_prefill_short_conv
+from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    report = {"cases": [], "maximum_row_relative_l2_limit": 0.001}
+    module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
+    torch.nn.Module.__init__(module)
+    module.conv_state_len = 9
+    module.short_conv_dilation = 3
+    for lengths in ([7344], [8976], [1600] * 4, [4, 0, 13, 7]):
+        n, rows, hidden = len(lengths), sum(lengths), 10240
+        torch.manual_seed(9683 + rows)
+        x = torch.randn(rows, hidden, device="cuda", dtype=torch.half) * 0.5
+        weight = torch.randn(hidden, 4, device="cuda", dtype=torch.half) * 0.125
+        seed = torch.randn(n + 2, hidden, 13, device="cuda", dtype=torch.half) * 0.5
+        control_state, candidate_state = seed.clone(), seed.clone()
+        starts = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum(
+            0, dtype=torch.int32
+        )
+        slots = torch.arange(1, n + 1, device="cuda", dtype=torch.int32)
+        initial = torch.arange(n, device="cuda") % 2 == 0
+        if lengths == [4, 0, 13, 7]:
+            slots[-1] = NULL_BLOCK_ID
+        metadata = SimpleNamespace(
+            non_spec_query_start_loc=starts,
+            has_initial_states_p=initial,
+            max_prefill_query_len=max(lengths),
+        )
+        max_length = max(lengths)
+
+        def control(
+            inputs=(x, metadata, control_state, weight, slots, n, 0, rows), seed=seed
+        ):
+            control_state = inputs[2]
+            control_state.copy_(seed)
+            return module._short_conv_dilated_prefill_batched(*inputs)
+
+        def candidate(
+            tile,
+            inputs=(
+                x,
+                candidate_state,
+                weight,
+                starts,
+                slots,
+                initial,
+                max_length,
+                0,
+                3,
+                NULL_BLOCK_ID,
+            ),
+            seed=seed,
+        ):
+            candidate_state = inputs[1]
+            candidate_state.copy_(seed)
+            return dilated_prefill_short_conv(
+                *inputs,
+                token_tile=tile,
+            )
+
+        def peak(fn):
+            torch.accelerator.synchronize()
+            torch.accelerator.reset_peak_memory_stats()
+            before = torch.accelerator.memory_allocated()
+            output = fn()
+            torch.accelerator.synchronize()
+            grown = torch.accelerator.max_memory_allocated() - before
+            del output
+            return grown
+
+        expected = control()
+        row = {
+            "lengths": lengths,
+            "control": measure(control, 5),
+            "control_peak_bytes": peak(control),
+            "candidates": [],
+        }
+        for tile in (4, 8, 16):
+            actual = candidate(tile)
+            error = relative_l2(actual, expected)
+            exact_state = torch.equal(candidate_state, control_state)
+            case = {
+                "token_tile": tile,
+                "paired_max_rel_l2": error,
+                "output_exact": torch.equal(actual, expected),
+                "state_exact": exact_state,
+                "passed": error <= 0.001
+                and exact_state
+                and bool(torch.isfinite(actual).all()),
+                **measure(lambda tile=tile: candidate(tile), 5),
+                "peak_bytes": peak(lambda tile=tile: candidate(tile)),
+            }
+            row["candidates"].append(case)
+        row["control_repeat"] = measure(control, 5)
+        report["cases"].append(row)
+        args.output.write_text(json.dumps(report, indent=2) + "\n")
+        print(json.dumps(row), flush=True)
+    raise SystemExit(
+        0 if all(c["passed"] for r in report["cases"] for c in r["candidates"]) else 1
+    )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/flashnext_prefill/qsa_prefill_tp.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TP4 paged-selector boundary/capture checks and complete-chain timings.
+
+The selector uses the serving PyNccl transport. All ranks compare output,
+causal/request isolation, caller-owned output and changed-input graph replay.
+This synthetic screen does not qualify model outputs or cache commit paths.
+"""
+
+import argparse
+import json
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+import torch.distributed as dist
+from collectives import measure
+
+from vllm.config import set_current_vllm_config
+from vllm.config.kernel import KernelConfig
+from vllm.config.sm70_sparse import Sm70SparseConfig
+from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+from vllm.models.qwen4_exp.nvidia.ops.qsa import qsa_select_paged_tokens
+
+
+class Group:
+    world_size = 4
+
+    def __init__(self, rank):
+        self.rank_in_group = rank
+        self.nccl = PyNcclCommunicator(dist.group.WORLD, rank)
+        self.calls = 0
+
+    def all_gather(self, tensor, dim):
+        assert dim == 0
+        self.calls += 1
+        out = tensor.new_empty((4 * tensor.shape[0], *tensor.shape[1:]))
+        self.nccl.all_gather(out, tensor)
+        return out
+
+
+def selection_error(actual, expected, positions):
+    a, e = actual.sort(1).values, expected.sort(1).values
+    locations = torch.searchsorted(e, a).clamp_max(e.shape[1] - 1)
+    replaced = (e.gather(1, locations) != a).sum(1)
+    positive = a >= 0
+    valid = bool(((a <= positions[:, None]) | ~positive).all())
+    distinct = bool(((a[:, 1:] > a[:, :-1]) | (a[:, 1:] == -1)).all())
+    same_padding = torch.equal((a == -1).sum(1), (e == -1).sum(1))
+    total = int(replaced.sum())
+    maximum = int(replaced.max()) if replaced.numel() else 0
+    return {
+        "exact": torch.equal(actual, expected),
+        "replaced_tokens": total,
+        "max_replacements_per_row": maximum,
+        "valid_distinct_causal": valid and distinct and same_padding,
+        "passed": valid
+        and distinct
+        and same_padding
+        and maximum <= 4
+        and total * 100000 <= actual.numel(),
+    }
+
+
+def run_case(group, rows, columns, *, mixed=False, graph=False):
+    torch.manual_seed(9683 + rows + columns)
+    page = 64
+    pages = (columns + page - 1) // page
+    cache = (
+        torch.randn(
+            pages * (2 if mixed else 1), page, 1, 128, device="cuda", dtype=torch.half
+        )
+        * 0.125
+    )
+    q = torch.randn(rows, 4, 128, device="cuda", dtype=torch.half) * 0.125
+    table = torch.arange(cache.shape[0], device="cuda", dtype=torch.int32).view(
+        -1, pages
+    )
+    # Retain real paged gathers and invalid-page masking.
+    table[:, 11::31] = -1
+    requests = torch.zeros(rows, device="cuda", dtype=torch.int32)
+    lengths = torch.full(
+        (table.shape[0],), columns * 4, device="cuda", dtype=torch.int32
+    )
+    positions = (
+        columns * 4 - rows + torch.arange(rows, device="cuda", dtype=torch.int32)
+    )
+    starts = None
+    if mixed:
+        requests[5:] = 1
+        positions[5:] -= 5
+        starts = torch.tensor([0, 5, rows], dtype=torch.int32)
+    expected_out = torch.empty(rows, 2051, device="cuda", dtype=torch.int32)
+    storage = torch.full((rows * 2051 + 2,), -77, device="cuda", dtype=torch.int32)
+    out = storage[1:-1].view(rows, 2051)
+
+    def run(tp, output):
+        return qsa_select_paged_tokens(
+            q,
+            cache,
+            table,
+            requests,
+            positions,
+            lengths,
+            2048,
+            4,
+            output,
+            query_start_loc_cpu=starts,
+            max_seq_len=columns * 4,
+            tp_group=tp,
+        )
+
+    expected = run(None, expected_out)
+    group.calls = 0
+    actual = run(group, out)
+    expected_calls = int(not mixed and 4096 <= rows <= 9216)
+    route = group.calls == expected_calls
+    errors = selection_error(actual, expected, positions)
+    errors["guard_exact"] = bool((storage[[0, -1]] == -77).all())
+    errors["output_identity"] = actual.data_ptr() == out.data_ptr()
+    errors["route_passed"] = route
+    errors["all_gather_calls"] = group.calls
+    errors["passed"] &= route and errors["guard_exact"] and errors["output_identity"]
+    row = {"rows": rows, "columns": columns, "mixed": mixed, "ranks": [None] * 4}
+    dist.all_gather_object(row["ranks"], errors)
+    if rows in (7344, 7345, 9216) and not mixed:
+        row["replicated"] = measure(lambda: run(None, expected_out))
+        row["partitioned"] = measure(lambda: run(group, out))
+        row["replicated_repeat"] = measure(lambda: run(None, expected_out))
+    if graph:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        captured = torch.cuda.CUDAGraph()
+        group.calls = 0
+        with torch.cuda.graph(captured, stream=stream):
+            run(group, out)
+        torch.cuda.current_stream().wait_stream(stream)
+        assert group.calls == 0
+        q.mul_(-0.5)
+        cache.add_(0.03125)
+        expected = run(None, expected_out)
+        captured.replay()
+        graph_errors = selection_error(out, expected, positions)
+        row["graph_ranks"] = [None] * 4
+        dist.all_gather_object(row["graph_ranks"], graph_errors)
+    row["passed"] = all(x["passed"] for x in row["ranks"] + row.get("graph_ranks", []))
+    return row
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    rank = int(os.environ["LOCAL_RANK"])
+    torch.accelerator.set_device_index(rank)
+    dist.init_process_group("gloo")
+    assert dist.get_world_size() == 4
+    group = Group(rank)
+    policy = Sm70SparseConfig(qsa_score_tile_mb=32)
+    policy.resolve()
+    config = SimpleNamespace(kernel_config=KernelConfig(sm70_sparse=policy))
+    report = {"cases": []}
+    try:
+        with set_current_vllm_config(config):
+            for rows, columns in [
+                (0, 1024),
+                (5, 1024),
+                (4095, 2048),
+                (4096, 2048),
+                (4097, 2048),
+                (7344, 1836),
+                (7345, 8672),
+                (9216, 8672),
+                (9217, 8672),
+                (4096, 65536),
+            ]:
+                row = run_case(group, rows, columns, graph=rows == 4097)
+                report["cases"].append(row)
+                if rank == 0:
+                    args.output.write_text(json.dumps(report, indent=2) + "\n")
+                    print(json.dumps(row), flush=True)
+            report["cases"].append(run_case(group, 4101, 2048, mixed=True))
+            if rank == 0:
+                args.output.write_text(json.dumps(report, indent=2) + "\n")
+    finally:
+        torch.accelerator.synchronize()
+        group.nccl.destroy()
+        dist.destroy_process_group()
+    raise SystemExit(0 if all(r["passed"] for r in report["cases"]) else 1)
+
+
+if __name__ == "__main__":
+    main()

docs/design/architecture/runtime_reference.md

@@ -24,10 +24,10 @@ identify initialization/hash code; they do not evaluate its semantics.
 | `vllm.config.diagnostic_dump.TensorDumpConfig` | kernel_config.sm70_moe.awq.diagnostics.compare_policy, kernel_config.sm70_moe.awq.diagnostics.dump_policy, kernel_config.sm70_moe.fp8.diagnostics.compare_policy, kernel_config.sm70_moe.fp8.diagnostics.dump_policy, observability_config.runtime_trace.dflash.pp_aux, observability_config.runtime_trace.dflash.tensors, observability_config.runtime_trace.dumps.awq_buffers, observability_config.runtime_trace.dumps.awq_compare, observability_config.runtime_trace.dumps.compile_inputs, observability_config.runtime_trace.dumps.fp8_compare, observability_config.runtime_trace.dumps.gdn_compare, observability_config.runtime_trace.dumps.gdn_core, observability_config.runtime_trace.dumps.gdn_graph, observability_config.runtime_trace.dumps.gdn_projection, observability_config.runtime_trace.dumps.mtp_step, observability_config.runtime_trace.dumps.qsa_calibration, observability_config.runtime_trace.dumps.qwen_layer, observability_config.runtime_trace.dumps.sample, observability_config.runtime_trace.dumps.sample_sync, observability_config.runtime_trace.dumps.sampler_logits, observability_config.runtime_trace.dumps.top1_sync, observability_config.runtime_trace.dumps.top_token_margin | resolve: [vllm/config/diagnostic_dump.py](../../../vllm/config/diagnostic_dump.py#L121); value: [vllm/config/diagnostic_dump.py](../../../vllm/config/diagnostic_dump.py#L243) | [vllm/config/diagnostic_dump.py](../../../vllm/config/diagnostic_dump.py#L77) |
 | `vllm.config.diagnostic_sampling.SamplingDiagnosticsConfig` | observability_config.runtime_trace.sampling | __post_init__: [vllm/config/diagnostic_sampling.py](../../../vllm/config/diagnostic_sampling.py#L55); value: [vllm/config/diagnostic_sampling.py](../../../vllm/config/diagnostic_sampling.py#L89) | [vllm/config/diagnostic_sampling.py](../../../vllm/config/diagnostic_sampling.py#L15) |
 | `vllm.config.ec_transfer.ECTransferConfig` | ec_transfer_config | compute_hash: [vllm/config/ec_transfer.py](../../../vllm/config/ec_transfer.py#L60); __post_init__: [vllm/config/ec_transfer.py](../../../vllm/config/ec_transfer.py#L78) | [vllm/config/ec_transfer.py](../../../vllm/config/ec_transfer.py#L16) |
-| `vllm.config.execution_policy.CommunicationPolicy` | parallel_config.communication, speculative_config.draft_parallel_config.communication, speculative_config.target_parallel_config.communication | value: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L553); resolve: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L558); compute_hash: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L588) | [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L544) |
+| `vllm.config.execution_policy.CommunicationPolicy` | parallel_config.communication, speculative_config.draft_parallel_config.communication, speculative_config.target_parallel_config.communication | value: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L561); resolve: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L566); compute_hash: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L596) | [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L552) |
 | `vllm.config.execution_policy.GraphPolicy` | compilation_config.runtime | resolve: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L100); compute_hash: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L170) | [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L88) |
-| `vllm.config.execution_policy.LayerExecutionPolicy` | kernel_config.layer_execution | value: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L336); resolve: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L341); compute_hash: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L380) | [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L276) |
-| `vllm.config.execution_policy.PlePlacementPolicy` | offload_config.ple | resolve: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L710); gib_bytes: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L724); compute_hash: [vllm/config/execution_policy_base.py](../../../vllm/config/execution_policy_base.py#L37) | [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L689) |
+| `vllm.config.execution_policy.LayerExecutionPolicy` | kernel_config.layer_execution | value: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L338); resolve: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L343); compute_hash: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L382) | [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L276) |
+| `vllm.config.execution_policy.PlePlacementPolicy` | offload_config.ple | resolve: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L718); gib_bytes: [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L732); compute_hash: [vllm/config/execution_policy_base.py](../../../vllm/config/execution_policy_base.py#L37) | [vllm/config/execution_policy.py](../../../vllm/config/execution_policy.py#L697) |
 | `vllm.config.flash_v100.FlashV100Diagnostics` | observability_config.runtime_trace.flash_v100 | resolve: [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L1311); value: [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L460); compile_ignored_aliases: [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L475) | [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L1303) |
 | `vllm.config.flash_v100.FlashV100Options` | attention_config.flash_v100.options | resolve: [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L504); value: [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L460); compile_ignored_aliases: [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L475) | [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L489) |
 | `vllm.config.flash_v100.FlashV100Policy` | attention_config.flash_v100 | bind_consumers: [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L1474); compile_ignored_aliases: [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L1478); resolve: [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L1511); compute_hash: [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L1516) | [vllm/config/flash_v100.py](../../../vllm/config/flash_v100.py#L1471) |
@@ -81,7 +81,7 @@ identify initialization/hash code; they do not evaluate its semantics.
 | `vllm.config.sm70_runtime.Sm70RuntimeConfig` | kernel_config.sm70_runtime | value: [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L89); __post_init__: [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L94); compile_ignored_aliases: [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L109) | [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L42) |
 | `vllm.config.sm70_runtime.SpecDecodeTraceConfig` | observability_config.spec_decode_trace | __post_init__: [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L424); resolve: [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L449); compile_ignored_aliases: [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L454) | [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L405) |
 | `vllm.config.sm70_runtime.StepProfilerConfig` | observability_config.step_profiler | __post_init__: [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L145); compile_ignored_aliases: [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L151) | [vllm/config/sm70_runtime.py](../../../vllm/config/sm70_runtime.py#L130) |
-| `vllm.config.sm70_sparse.Sm70SparseConfig` | kernel_config.sm70_sparse | compute_hash: [vllm/config/sm70_sparse.py](../../../vllm/config/sm70_sparse.py#L166); resolve: [vllm/config/execution_policy_base.py](../../../vllm/config/execution_policy_base.py#L57); value: [vllm/config/execution_policy_base.py](../../../vllm/config/execution_policy_base.py#L69) | [vllm/config/sm70_sparse.py](../../../vllm/config/sm70_sparse.py#L51) |
+| `vllm.config.sm70_sparse.Sm70SparseConfig` | kernel_config.sm70_sparse | compute_hash: [vllm/config/sm70_sparse.py](../../../vllm/config/sm70_sparse.py#L168); resolve: [vllm/config/execution_policy_base.py](../../../vllm/config/execution_policy_base.py#L57); value: [vllm/config/execution_policy_base.py](../../../vllm/config/execution_policy_base.py#L69) | [vllm/config/sm70_sparse.py](../../../vllm/config/sm70_sparse.py#L51) |
 | `vllm.config.sm70_triton_attention.Sm70TritonAttentionPolicy` | attention_config.sm70_triton | resolve: [vllm/config/sm70_triton_attention.py](../../../vllm/config/sm70_triton_attention.py#L73); compute_hash: [vllm/config/sm70_triton_attention.py](../../../vllm/config/sm70_triton_attention.py#L110); value: [vllm/config/execution_policy_base.py](../../../vllm/config/execution_policy_base.py#L69) | [vllm/config/sm70_triton_attention.py](../../../vllm/config/sm70_triton_attention.py#L36) |
 | `vllm.config.speculative.SpeculativeConfig` | speculative_config | compute_hash: [vllm/config/speculative.py](../../../vllm/config/speculative.py#L398); __post_init__: [vllm/config/speculative.py](../../../vllm/config/speculative.py#L740) | [vllm/config/speculative.py](../../../vllm/config/speculative.py#L128) |
 | `vllm.config.speculative_sampling.SpeculativeSamplingPolicy` | speculative_config.sampling_policy | resolve_fields: [vllm/config/speculative_sampling.py](../../../vllm/config/speculative_sampling.py#L109); resolve: [vllm/config/speculative_sampling.py](../../../vllm/config/speculative_sampling.py#L142); compute_hash: [vllm/config/speculative_sampling.py](../../../vllm/config/speculative_sampling.py#L181) | [vllm/config/speculative_sampling.py](../../../vllm/config/speculative_sampling.py#L16) |

docs/design/flashnext_indexer_placement_20261015.md

@@ -177,3 +177,56 @@ FP16 draft HOT16384, 80/20 expert placement, 256K context and chunk8192.
 Whole-model peak live/driver memory, verifier distribution, cold prefill,
 long-context quality and fresh-process reproduction remain pending; no
 512-MiB headroom or 16-ms verifier acceptance is claimed from operator tests.
+
+## Ordinary verifier window and peer prefill port
+
+The standard Worker restart, without step-event fences, passes the matched
+32K/64K/128K sweep at 152.15 / 139.99 / 110.49 output tok/s. Natural-EOS
+quality again passes through 260977 tokens, taking 113.85 seconds at the long
+boundary. The equivalent preceding single-stage diagnostic run took114.02s,
+compared with345.94s on the 128K-stage workaround; full staging repairs most of
+that cliff. A diagnostic private CPU-pool screen gives 15.96ms rank-max verifier
+mean at32K (p99 23.97ms); this is kept separate from ordinary tracing.
+
+Ordinary CUDA graph-envelope tracing measures504 committed target rounds over
+eight matched192-token32K requests. Every rank has512 state-commit triples:
+one prefill plus63 metric-reconciled target verifications per request. There
+is no extra pipeline-end commit in this trace. All504 target graph starts are
+present and associate uniquely with their following same-stream commit. This
+is distinguished from the older trace's extra pipeline-end replay; the parser
+has an explicit tail-count selection, rather than imputing or dropping events.
+
+From the earliest target GPU start to the latest target state-commit end across
+TP4, mean is **15.216741ms**, p50 14.5569745ms, p90 16.915299ms, p99 24.133488ms,
+maximum40.109622ms. All504 windows are exact; there are no conservative missing
+start bounds. This passes the historical mean<=16ms criterion, not an every
+round or p99 target. The source is996e3a161c, main history on host, compressed
+indexing on GPU, MTP4/FULL M5, encoded HOT32768 and FP16 draft HOT16384. Helper,
+EngineCore and PLE threads use disjoint physical-core pools. Standard NCCL
+MAX_CTAS6/buffer4MiB uses390070272 bytes in the two-group screen,48MiB less
+than8CTAs/buffer4MiB, with0.499ms all-gather. Unprofiled output throughput,
+round estimates and this independently traced verifier window remain separate.
+
+Further reuse of peer commit31edb1fe3d adds only its qualified QSA TP-row and PLE
+prefill stages. Target/decode graphs and native sources remain unchanged. The
+QSA adapter qualifies replicated queries/history for TP4/DCP1/PCP1; the selector
+admits single-request FP16/H4/D128 prefill at4096..9216 query rows with sufficient
+score work, gathers compressed512-block selections and expands globally.
+Decode, capture, mixed requests, threshold neighbors and other layouts retain
+the existing selector. The per-tile logits lifetime repair stays inside the
+loop for both paths. The opt-in remains a typed active-QSA computation policy.
+
+The PLE stage reuses the generic PLE kernel owner, consumes existing input and
+state directly, commits valid base state separately and retains the native FP16
+SiLU boundary. SM70/FP16/10240channels/4taps/dilation3/base9 are the accelerated
+subset; unsupported layouts use the existing convolution. Both opt-ins default
+to false in the library. No private kernel artifact or native ABI update is
+required: the Triton source is shipped with the ordinary package.
+
+Local port checks:64CPU cases pass with24GPU skips;50 PLE GPU cases pass. At the
+actual7344-row geometry, compared with this branch's bounded quarter-channel
+control, fused PLE reduces additional peak from227123712 to150405120 bytes,
+with bitwise-equal output and state, and0.85545 versus7.34761ms in the repeated
+control. Destination affected suites pass98cases. Destination TP4 selector and
+final combined serving/quality/memory gates are pending; neither opt-in is
+promoted from these operator checks alone.

tests/config/test_sm70_provider_lifecycle.py

@@ -262,6 +262,34 @@ def test_sparse_hash_and_errors_only_include_consumed_family(
         policy.validate_active()
 
 
+@pytest.mark.parametrize("family", ["qsa", "indexer"])
+def test_prefill_row_partition_is_hashed_for_its_family(family):
+    first = Sm70SparseConfig()
+    second = Sm70SparseConfig(qsa_prefill_tp_rows=True)
+    for policy in (first, second):
+        policy.resolve()
+        policy.qualify(family)
+    assert first.compute_hash() == second.compute_hash()  # Inactive owner.
+    first.active = second.active = True
+    assert (first.compute_hash() != second.compute_hash()) == (family == "qsa")
+    second.qsa_prefill_tp_rows = False
+    assert first.compute_hash() == second.compute_hash()
+
+
+def test_prefill_convolution_policy_is_engine_local_and_hashed():
+    from vllm.config.kernel import KernelConfig
+
+    first, second = KernelConfig(), KernelConfig()
+    first.layer_execution.resolve()
+    second.layer_execution.resolve()
+    baseline = first.layer_execution.compute_hash()
+    second.layer_execution.ple_prefill_conv = True
+    assert first.layer_execution.compute_hash() == baseline
+    assert second.layer_execution.compute_hash() != baseline
+    second.layer_execution.ple_prefill_conv = False
+    assert second.layer_execution.compute_hash() == baseline
+
+
 @pytest.mark.parametrize(
     "raw,expected", [(None, True), ("", False), ("2", False), ("1", True)]
 )

tests/models/qwen4_exp/test_ple_short_conv_prefill.py

@@ -469,3 +469,77 @@ def test_long_single_prefill_channel_chunks_release_device_headroom(monkeypatch)
     assert chunked <= 1.9 * batch_bytes, (chunked, batch_bytes)
     assert single - chunked >= 0.4 * batch_bytes, (single, chunked, batch_bytes)
     print(f"single_prefill_peak_bytes: single={single} chunked={chunked}")
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("length", [0, 1, 8, 9, 10, 16, 17, 2049])
+@pytest.mark.parametrize("initial", [False, True])
+def test_sm70_prefill_convolution_output_and_state(length, initial):
+    if torch.cuda.get_device_capability() != (7, 0):
+        pytest.skip("SM70 admission")
+    module, x, args = _case(
+        "cuda",
+        torch.float16,
+        lengths=[length, 0, 3],
+        decode_tokens=2,
+        state_indices=[3, 4, NULL_BLOCK_ID],
+        has_initial=[initial, True, True],
+        hidden_size=10240,
+    )
+    metadata, state, *rest = args
+    expected_state = state.clone()
+    expected = module._short_conv_dilated_prefill_batched(
+        x, metadata, expected_state, *rest
+    )
+    module._prefill_short_conv = True
+    actual = module._short_conv_dilated_prefill_batched(x, metadata, state, *rest)
+    assert torch.equal(state, expected_state)
+    assert bool(torch.isfinite(actual).all())
+    # Preserve the operation-aware predeclared FP16 bound. These constructed
+    # cases also record the stronger exact equality through assert_close.
+    error = (actual.float() - expected.float()).norm(dim=1)
+    scale = expected.float().norm(dim=1).clamp_min(1e-6)
+    assert (error / scale).max() <= 0.001
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("transpose_state", [False, True])
+def test_sm70_prefill_convolution_changed_input_graph_and_continuation(transpose_state):
+    if torch.cuda.get_device_capability() != (7, 0):
+        pytest.skip("SM70 admission")
+    module, x, args = _case(
+        "cuda",
+        torch.float16,
+        lengths=[7, 18],
+        decode_tokens=1,
+        state_indices=[4, 2],
+        has_initial=[True, False],
+        hidden_size=10240,
+    )
+    metadata, state, *rest = args
+    if transpose_state:
+        state = state.transpose(1, 2).contiguous().transpose(1, 2)
+    seed = state.clone()
+    module._prefill_short_conv = True
+    module._short_conv_dilated_prefill_batched(x, metadata, state, *rest)
+    state.copy_(seed)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        actual = module._short_conv_dilated_prefill_batched(x, metadata, state, *rest)
+    torch.cuda.current_stream().wait_stream(stream)
+    state.copy_(seed)
+    x.mul_(-0.75)
+    metadata.has_initial_states_p.fill_(True)
+    rest[1].copy_(rest[1].flip(0))
+    reference_state = state.clone()
+    module._prefill_short_conv = False
+    for _ in range(3):
+        expected = module._short_conv_dilated_prefill_batched(
+            x, metadata, reference_state, *rest
+        )
+        graph.replay()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        assert torch.equal(state, reference_state)

vllm/config/execution_policy.py

@@ -282,6 +282,8 @@ class LayerExecutionPolicy(ExecutionPolicy):
     """Retain the independent SM70 shared-expert overlap rollback."""
     ple_spec_conv: bool | None = None
     """Retain the existing PLE MTP convolution and state-commit kernel."""
+    ple_prefill_conv: bool = False
+    """Use the FP16 PLE prefill convolution without a padded history."""
     quant_backend: Literal["auto", "marlin", "turbomind"] | None = None
     """Shared pre-Ampere quantization backend; format flags retain their gates."""
     gptq_turbomind: bool | None = None
@@ -379,6 +381,12 @@ class LayerExecutionPolicy(ExecutionPolicy):
 
     def compute_hash(self) -> str:
         factors: dict[str, object] = {"policy": super().compute_hash()}
+        if (
+            self.active
+            and self.ple_prefill_conv
+            and (self.hash_fields is None or "ple_spec_conv" in self.hash_fields)
+        ):
+            factors["ple_prefill_conv"] = True
         errors = {
             field: error
             for field, error in self.provider_errors.items()

vllm/config/sm70_sparse.py

@@ -77,6 +77,8 @@ class Sm70SparseConfig(DeferredExecutionPolicy):
     """Retain the M5/M10 sparse top-k compaction route."""
     qsa_score_tile_mb: int | None = None
     """Score workspace budget; retain the qualified 64 MiB default."""
+    qsa_prefill_tp_rows: bool = False
+    """Partition eligible replicated QSA prefill scoring across four TP ranks."""
     qsa_cublas_min_rows: int | None = None
     """Minimum query rows for the cuBLAS scorer."""
     qsa_cublas_min_score_elements: int | None = None
@@ -167,6 +169,11 @@ class Sm70SparseConfig(DeferredExecutionPolicy):
         if not self.active:
             return hash_factors({})
         factors: dict[str, object] = {"policy": super().compute_hash()}
+        if self.qsa_prefill_tp_rows and (
+            self.hash_fields is None
+            or any(field.startswith("qsa_") for field in self.hash_fields)
+        ):
+            factors["qsa_prefill_tp_rows"] = True
         if self.hash_fields is None or any(
             not field.startswith("qsa_") for field in self.hash_fields
         ):

vllm/model_executor/kernels/ple/short_conv.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dilated PLE prefill convolution without a padded channels-first history."""
+
+import torch
+import torch.nn.functional as F
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _dilated_prefill_conv(
+    X,
+    State,
+    Weight,
+    Starts,
+    Slots,
+    Initial,
+    Out,
+    D: tl.constexpr,
+    X_STRIDE: tl.constexpr,
+    S0: tl.constexpr,
+    S1: tl.constexpr,
+    S2: tl.constexpr,
+    W0: tl.constexpr,
+    W1: tl.constexpr,
+    DECODE_TOKENS: tl.constexpr,
+    STATE_AVAILABLE: tl.constexpr,
+    NULL_SLOT: tl.constexpr,
+    K: tl.constexpr,
+    DILATION: tl.constexpr,
+    BT: tl.constexpr,
+    BD: tl.constexpr,
+):
+    req = tl.program_id(2)
+    first = tl.load(Starts + req) - DECODE_TOKENS
+    length = tl.load(Starts + req + 1) - DECODE_TOKENS - first
+    slot = tl.load(Slots + req).to(tl.int64)
+    valid = slot != NULL_SLOT
+    initial = tl.load(Initial + req).to(tl.int1) & valid
+    times = tl.program_id(0) * BT + tl.arange(0, BT)
+    cols = tl.program_id(1) * BD + tl.arange(0, BD)
+    active = (times[:, None] < length) & (cols[None, :] < D)
+    acc = tl.full((BT, BD), 0.0, tl.float32)
+    for tap in tl.static_range(K):
+        pos = times - (K - 1 - tap) * DILATION
+        value = tl.load(
+            X + (first + pos[:, None]) * X_STRIDE + cols[None, :],
+            active & (pos[:, None] >= 0) & valid,
+            other=0.0,
+        ).to(tl.float32)
+        if STATE_AVAILABLE:
+            old = tl.load(
+                State
+                + slot * S0
+                + cols[None, :] * S1
+                + ((K - 1) * DILATION + pos[:, None]) * S2,
+                active & (pos[:, None] < 0) & initial,
+                other=0.0,
+            ).to(tl.float32)
+            value += old
+        weight = tl.load(Weight + cols * W0 + tap * W1, cols < D, 0).to(tl.float32)
+        acc = tl.fma(value, weight[None, :], acc)
+    tl.store(Out + (first + times[:, None]) * D + cols[None, :], acc, active)
+
+
+@triton.jit
+def _dilated_prefill_state(
+    X,
+    State,
+    Starts,
+    Slots,
+    Initial,
+    D: tl.constexpr,
+    X_STRIDE: tl.constexpr,
+    S0: tl.constexpr,
+    S1: tl.constexpr,
+    S2: tl.constexpr,
+    DECODE_TOKENS: tl.constexpr,
+    NULL_SLOT: tl.constexpr,
+    LENGTH: tl.constexpr,
+    BL: tl.constexpr,
+    BD: tl.constexpr,
+):
+    req = tl.program_id(1)
+    first = tl.load(Starts + req) - DECODE_TOKENS
+    length = tl.load(Starts + req + 1) - DECODE_TOKENS - first
+    slot = tl.load(Slots + req).to(tl.int64)
+    valid = (slot != NULL_SLOT) & (length > 0)
+    initial = tl.load(Initial + req).to(tl.int1)
+    cols = tl.program_id(0) * BD + tl.arange(0, BD)
+    offsets = tl.arange(0, BL)
+    pos = length - LENGTH + offsets
+    mask = (cols[:, None] < D) & (offsets[None, :] < LENGTH) & valid
+    recent = tl.load(
+        X + (first + pos[None, :]) * X_STRIDE + cols[:, None],
+        mask & (pos[None, :] >= 0),
+        other=0.0,
+    )
+    # One CTA owns all history positions for a channel. Read the old window
+    # before its in-place update; other CTAs never write these channels.
+    old = tl.load(
+        State + slot * S0 + cols[:, None] * S1 + (LENGTH + pos[None, :]) * S2,
+        mask & (pos[None, :] < 0) & initial,
+        other=0.0,
+    )
+    next_state = tl.where(pos[None, :] >= 0, recent, old)
+    tl.store(
+        State + slot * S0 + cols[:, None] * S1 + offsets[None, :] * S2,
+        next_state,
+        mask,
+    )
+
+
+def dilated_prefill_short_conv(
+    x,
+    state,
+    weight,
+    starts,
+    slots,
+    initial,
+    max_length,
+    decode_tokens,
+    dilation,
+    null_slot,
+    *,
+    token_tile=8,
+    channel_tile=128,
+):
+    """Compute request-local convolution, then commit each valid base state.
+
+    Callers provide contiguous channels, unique non-null state slots, and the
+    prefill suffix of query starts. Extra speculative state positions remain
+    unchanged. This stage is screened independently before runtime admission.
+    """
+    rows, hidden = x.shape
+    out = torch.empty_like(x, memory_format=torch.contiguous_format)
+    if rows == 0:
+        return out
+    kernel = weight.shape[1]
+    common = dict(
+        D=hidden,
+        X_STRIDE=x.stride(0),
+        S0=state.stride(0),
+        S1=state.stride(1),
+        S2=state.stride(2),
+        DECODE_TOKENS=decode_tokens,
+        NULL_SLOT=null_slot,
+    )
+    _dilated_prefill_conv[
+        (
+            triton.cdiv(max_length, token_tile),
+            triton.cdiv(hidden, channel_tile),
+            slots.numel(),
+        )
+    ](
+        x,
+        state,
+        weight,
+        starts,
+        slots,
+        initial,
+        out,
+        **common,
+        W0=weight.stride(0),
+        W1=weight.stride(1),
+        STATE_AVAILABLE=state.shape[0] > 0,
+        K=kernel,
+        DILATION=dilation,
+        BT=token_tile,
+        BD=channel_tile,
+        num_warps=4,
+    )
+    state_length = (kernel - 1) * dilation
+    if state_length and state.shape[0]:
+        _dilated_prefill_state[(triton.cdiv(hidden, 64), slots.numel())](
+            x,
+            state,
+            starts,
+            slots,
+            initial,
+            **common,
+            LENGTH=state_length,
+            BL=triton.next_power_of_2(state_length),
+            BD=64,
+            num_warps=4,
+        )
+    # Retain the native FP16 conv -> SiLU rounding boundary.
+    return F.silu(out, inplace=True)

vllm/models/qwen4_exp/nvidia/indexer_qsa.py

@@ -130,6 +130,17 @@ class QSAIndexer(nn.Module):
         # while continuing to update the QSA side cache.
         self.skip_topk = False
         self.shared_key_scoring = vllm_config.kernel_config.sm70_qsa_shared_key
+        self.prefill_tp_group = None
+        parallel = vllm_config.parallel_config
+        if (
+            vllm_config.kernel_config.sm70_sparse.qsa_prefill_tp_rows
+            and parallel.tensor_parallel_size == 4
+            and parallel.decode_context_parallel_size == 1
+            and parallel.prefill_context_parallel_size == 1
+        ):
+            from vllm.distributed.parallel_state import get_tp_group
+
+            self.prefill_tp_group = get_tp_group()
 
         self.index_qk_proj = ReplicatedLinear(
             int(config.hidden_size),
@@ -362,6 +373,7 @@ class QSAIndexer(nn.Module):
             query_start_loc_cpu=compressed_metadata.query_start_loc_cpu,
             shared_key_scoring=self.shared_key_scoring,
             max_seq_len=compressed_metadata.max_seq_len,
+            tp_group=self.prefill_tp_group,
         )
 
 

vllm/models/qwen4_exp/nvidia/ops/qsa.py

@@ -1600,6 +1600,7 @@ def qsa_select_paged_tokens(
     query_start_loc_cpu: torch.Tensor | None = None,
     shared_key_scoring: bool = False,
     max_seq_len: int | None = None,
+    tp_group=None,
 ) -> torch.Tensor:
     """Score, select, and expand QSA indices without host synchronization.
 
@@ -1670,6 +1671,35 @@ def qsa_select_paged_tokens(
         else:
             score_columns = capacity_columns
             all_visible = None
+    # Only the model adapter may qualify replicated queries/cache for TP.
+    # Keep small work, mixed requests, other topologies and graph capture on
+    # the existing selector. Cache writes have already completed on every rank.
+    tp_step = 0
+    if (
+        tp_group is not None
+        and tp_group.world_size == 4
+        and 4096 <= rows <= 9216
+        and q.shape[1:] == (4, 128)
+        and token_topk == 2048
+        and compress_ratio == 4
+        and contiguous_keys is not None
+        and not torch.cuda.is_current_stream_capturing()
+    ):
+        step = triton.cdiv(rows, tp_group.world_size)
+        smallest_partition = rows - (tp_group.world_size - 1) * step
+        if smallest_partition >= sparse_policy().value(
+            "qsa_cublas_min_rows"
+        ) and _qsa_indexer_cublas_work_supported(smallest_partition, score_columns):
+            tp_step = step
+            first = tp_group.rank_in_group * step
+            last = min(first + step, rows)
+            q = q[first:last]
+            assert all_visible is not None
+            all_visible = all_visible[first:last]
+            rows = last - first
+            logger.info_once(
+                "Using SM70 QSA TP4 prefill row partition with block-index AllGather."
+            )
     # Bound the per-chunk logits by the configured score tile on both the
     # cuBLAS and paged paths. Startup profiling skips attention, so this
     # transient is not in the KV budget; the default 64 MiB tile preserves the
@@ -1681,7 +1711,7 @@ def qsa_select_paged_tokens(
     rows_per_chunk = max(1, logits_bytes // max(score_columns * 4, 1))
     chunk_rows = min(rows, rows_per_chunk)
     blocks_buffer = torch.empty(
-        (chunk_rows, block_topk), dtype=torch.int32, device=q.device
+        (tp_step or chunk_rows, block_topk), dtype=torch.int32, device=q.device
     )
     exact_topk = _use_sm70_qsa_lexicographic_topk(block_topk)
     topk_workspace = (
@@ -1711,7 +1741,11 @@ def qsa_select_paged_tokens(
                 compress_ratio,
                 shared_key_scoring=shared_key_scoring,
             )
-        blocks = blocks_buffer[: row_end - row_start]
+        blocks = (
+            blocks_buffer[row_slice]
+            if tp_step
+            else blocks_buffer[: row_end - row_start]
+        )
         use_cooperative_topk = (
             blocks.shape[0] <= 32
             and logits.stride(0) % 4 == 0
@@ -1746,18 +1780,32 @@ def qsa_select_paged_tokens(
                 block_topk,
                 score_columns,
             )
+        if not tp_step:
+            expand_qsa_block_indices_cuda(
+                blocks,
+                query_positions[row_slice],
+                sequence_lengths,
+                token_to_req[row_slice],
+                compress_ratio,
+                token_topk,
+                out[row_slice],
+            )
+        # Each rank has consumed this tile, including in the TP row partition.
+        del logits, visible_blocks
+    if tp_step:
+        # Gather the compressed block IDs, not the four-times wider expanded
+        # token matrix. Padding is confined to the last rank and discarded.
+        blocks_buffer[rows:].fill_(-1)
+        gathered = tp_group.all_gather(blocks_buffer, dim=0)
         expand_qsa_block_indices_cuda(
-            blocks,
-            query_positions[row_slice],
+            gathered[: out.shape[0]],
+            query_positions,
             sequence_lengths,
-            token_to_req[row_slice],
+            token_to_req,
             compress_ratio,
             token_topk,
-            out[row_slice],
+            out,
         )
-        # Selection and expansion have consumed the scores on the same stream.
-        # Avoid overlapping two full logits tiles on the next iteration.
-        del logits, visible_blocks
     return out
 
 

vllm/models/qwen4_exp/nvidia/ple_layer.py

@@ -2175,6 +2175,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
         self.conv_kernel_size = int(config.ple_conv_kernel_size)
         self.short_conv_dilation = int(config.ngram_size)
         self.conv_state_len = (self.conv_kernel_size - 1) * self.short_conv_dilation
+        self._prefill_short_conv = (
+            vllm_config.kernel_config.layer_execution.ple_prefill_conv
+        )
         self.num_spec_tokens = vllm_config.num_speculative_tokens
         self.activation = "silu"
         # The offload subprocess constructs the surrounding model on meta, but
@@ -2526,6 +2529,54 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
 
         hidden_size = x_p.shape[1]
         state_len = self.conv_state_len
+        if (
+            getattr(self, "_prefill_short_conv", False)
+            and current_platform.is_device_capability(70)
+            and hidden_size == 10240
+            and state_len == 9
+            and self.short_conv_dilation == 3
+            and conv_weights.shape == (hidden_size, 4)
+            and conv_state.ndim == 3
+            and conv_state.shape[1] == hidden_size
+            and conv_state.shape[2] >= state_len
+            and x_p.stride(1) == 1
+            and all(
+                t.is_cuda and t.dtype == torch.float16
+                for t in (x_p, conv_state, conv_weights)
+            )
+            and all(
+                t.device == x_p.device
+                for t in (
+                    conv_state,
+                    conv_weights,
+                    q_starts,
+                    state_indices_tensor_p,
+                    has_initial_states_p,
+                )
+            )
+            and state_indices_tensor_p.is_contiguous()
+            and has_initial_states_p.is_contiguous()
+        ):
+            from vllm.model_executor.kernels.ple.short_conv import (
+                dilated_prefill_short_conv,
+            )
+
+            logger.info_once(
+                "Using SM70 PLE prefill convolution without padded history."
+            )
+            return dilated_prefill_short_conv(
+                x_p,
+                conv_state,
+                conv_weights,
+                q_starts,
+                state_indices_tensor_p[:num_prefills],
+                has_initial_states_p[:num_prefills],
+                max_len,
+                0,
+                self.short_conv_dilation,
+                NULL_BLOCK_ID,
+                token_tile=16,
+            )
         positions = torch.arange(
             num_prefill_tokens, device=x_p.device, dtype=torch.int64
         )