diff --git a/benchmarks/flashnext_prefill/ple_prefill_conv.py b/benchmarks/flashnext_prefill/ple_prefill_conv.py new file mode 100644 index 000000000..70db24e49 --- /dev/null +++ b/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() diff --git a/benchmarks/flashnext_prefill/qsa_prefill_tp.py b/benchmarks/flashnext_prefill/qsa_prefill_tp.py new file mode 100644 index 000000000..ca58b3af1 --- /dev/null +++ b/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() diff --git a/docs/design/architecture/runtime_reference.md b/docs/design/architecture/runtime_reference.md index 31ab309db..c56aa0336 100644 --- a/docs/design/architecture/runtime_reference.md +++ b/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) | diff --git a/docs/design/flashnext_indexer_placement_20261015.md b/docs/design/flashnext_indexer_placement_20261015.md index 6b93f5500..de884947c 100644 --- a/docs/design/flashnext_indexer_placement_20261015.md +++ b/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. diff --git a/tests/config/test_sm70_provider_lifecycle.py b/tests/config/test_sm70_provider_lifecycle.py index adf28eb7d..a0835c427 100644 --- a/tests/config/test_sm70_provider_lifecycle.py +++ b/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)] ) diff --git a/tests/models/qwen4_exp/test_ple_short_conv_prefill.py b/tests/models/qwen4_exp/test_ple_short_conv_prefill.py index f0a31f037..b2bf2a0c5 100644 --- a/tests/models/qwen4_exp/test_ple_short_conv_prefill.py +++ b/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) diff --git a/vllm/config/execution_policy.py b/vllm/config/execution_policy.py index 59ba83d44..a29d3d522 100644 --- a/vllm/config/execution_policy.py +++ b/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() diff --git a/vllm/config/sm70_sparse.py b/vllm/config/sm70_sparse.py index a5327177f..923923560 100644 --- a/vllm/config/sm70_sparse.py +++ b/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 ): diff --git a/vllm/model_executor/kernels/ple/short_conv.py b/vllm/model_executor/kernels/ple/short_conv.py new file mode 100644 index 000000000..2e20ff979 --- /dev/null +++ b/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) diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py index ababa83d6..ad7ea2e44 100644 --- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py +++ b/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, ) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index b9a9a2202..8c272eecf 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/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 diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 074924026..e71000814 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/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 )