diff --git a/benchmarks/benchmark_infllmv2.py b/benchmarks/benchmark_infllmv2.py new file mode 100644 index 00000000000..6195a805f6d --- /dev/null +++ b/benchmarks/benchmark_infllmv2.py @@ -0,0 +1,2402 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run and compare traceable dense/InfLLM-V2 long-context benchmarks. + +The ``run`` command delegates request generation and latency collection to +``benchmark_serving.py`` while sampling total memory used on one GPU. The +``report`` command only compares run files whose workload fingerprints match. +The ``sparse-diagnostic`` command binds a newly generated selector trace to the +exact random-token workload without treating its trace-instrumented latency as +a performance result. A timed sparse run requires that bound diagnostic. +The ``operators`` command uses CUDA events to time compressed-K update, Stage 1, +Stage 2, and their complete chain for a configurable context/concurrency matrix. +The ``prefill`` command compares complete sparse and dense Paddle prefill paths. +The ``cuda-impl`` command runs the equivalent decode workload against a local +``infllm_v2`` PyTorch extension checkout. + +Selector hit rate is the micro-averaged recall of dense-reference blocks: + + sum(|selected_blocks & reference_blocks|) / sum(|reference_blocks|) + +Each selector sample represents one request/query/KV-head tuple. Unbound raw +samples and summaries cannot be attached to timed sparse results. +""" + +from __future__ import annotations + +import argparse +import csv +import datetime +import hashlib +import json +import math +import os +import statistics +import subprocess +import sys +import tempfile +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +SCHEMA_VERSION = 2 +BENCHMARK_SERVING = Path(__file__).with_name("benchmark_serving.py") +SELECTOR_DEFINITION = "micro recall of dense-reference blocks in selected blocks" +SELECTOR_REFERENCE_DEFINITION = ( + "top blocks by exact dense softmax attention mass, summed over the query heads sharing one KV head" +) +SELECTOR_TRACE_PATH_ENV = "FD_INFLLMV2_SELECTOR_TRACE_PATH" +STAGE2_BLOCKS_PER_SPLIT = 2 +SERVING_RESULT_FIELDS = ( + "duration", + "completed", + "total_input_tokens", + "total_output_tokens", + "request_throughput", + "output_throughput", + "total_token_throughput", + "mean_ttft_ms", + "median_ttft_ms", + "p99_ttft_ms", + "mean_tpot_ms", + "median_tpot_ms", + "p99_tpot_ms", +) + + +@dataclass(frozen=True) +class SelectorMetrics: + selector_hit_rate: float + selector_samples: int + selector_hits: int | None + selector_targets: int | None + definition: str = SELECTOR_DEFINITION + + +@dataclass(frozen=True) +class GPUMemoryMetrics: + gpu_index: int + gpu_uuid: str + gpu_name: str + total_mib: int + baseline_mib: int + peak_mib: int + peak_delta_mib: int + samples: int + sample_interval_seconds: float + + +@dataclass(frozen=True) +class GPUIdentity: + gpu_index: int + gpu_uuid: str + gpu_name: str + total_mib: int + + +@dataclass(frozen=True) +class OperatorScenario: + context_length: int + concurrency: int + + +def _read_json(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as json_file: + value = json.load(json_file) + if not isinstance(value, dict): + raise TypeError(f"{path} must contain a JSON object, got {type(value).__name__}.") + return value + + +def _sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as binary_file: + for chunk in iter(lambda: binary_file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _loaded_shared_object_metadata(filename: str) -> dict[str, Any]: + mapped_paths = set() + with Path("/proc/self/maps").open("r", encoding="utf-8") as maps_file: + for line in maps_file: + mapped_path = line.rstrip().split(maxsplit=5)[-1] + if mapped_path.endswith(filename): + mapped_paths.add(Path(mapped_path).resolve()) + if len(mapped_paths) != 1: + raise RuntimeError( + f"Expected exactly one loaded {filename}, found " f"{[str(path) for path in sorted(mapped_paths)]}." + ) + path = mapped_paths.pop() + stat_result = path.stat() + return { + "path": str(path), + "mtime_ns": stat_result.st_mtime_ns, + "mtime_utc": datetime.datetime.fromtimestamp(stat_result.st_mtime, datetime.timezone.utc).isoformat(), + "size_bytes": stat_result.st_size, + "sha256": _sha256_file(path), + } + + +def _read_json_snapshot(path: Path) -> tuple[dict[str, Any], str, os.stat_result]: + stat_before = path.stat() + if not path.is_file(): + raise ValueError(f"{path} must be a regular file.") + with path.open("rb") as binary_file: + raw_value = binary_file.read() + stat_after = path.stat() + if ( + stat_before.st_dev, + stat_before.st_ino, + stat_before.st_size, + stat_before.st_mtime_ns, + ) != ( + stat_after.st_dev, + stat_after.st_ino, + stat_after.st_size, + stat_after.st_mtime_ns, + ): + raise RuntimeError(f"{path} changed while it was being read.") + value = json.loads(raw_value) + if not isinstance(value, dict): + raise TypeError(f"{path} must contain a JSON object, got {type(value).__name__}.") + return value, _sha256_bytes(raw_value), stat_after + + +def _write_json(path: Path, value: dict[str, Any], overwrite: bool) -> None: + if not path.parent.is_dir(): + raise FileNotFoundError(f"Output directory does not exist: {path.parent}") + mode = "w" if overwrite else "x" + with path.open(mode, encoding="utf-8") as json_file: + json.dump(value, json_file, ensure_ascii=False, indent=2, sort_keys=True) + json_file.write("\n") + + +def _require_int(value: Any, name: str, minimum: int) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer, got {type(value).__name__}.") + if value < minimum: + raise ValueError(f"{name} must be >= {minimum}, got {value}.") + return value + + +def _require_number(value: Any, name: str, minimum: float, inclusive: bool = True) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"{name} must be a number, got {type(value).__name__}.") + number = float(value) + if not math.isfinite(number): + raise ValueError(f"{name} must be finite, got {number}.") + invalid = number < minimum if inclusive else number <= minimum + if invalid: + operator = ">=" if inclusive else ">" + raise ValueError(f"{name} must be {operator} {minimum}, got {number}.") + return number + + +def _require_rate(value: Any, name: str) -> float: + rate = _require_number(value, name, 0.0) + if rate > 1.0: + raise ValueError(f"{name} must be <= 1.0, got {rate}.") + return rate + + +def _require_sha256(value: Any, name: str) -> str: + if not isinstance(value, str) or len(value) != 64: + raise ValueError(f"{name} must be a lowercase SHA256 hex digest.") + if any(character not in "0123456789abcdef" for character in value): + raise ValueError(f"{name} must be a lowercase SHA256 hex digest.") + return value + + +def _validate_block_ids(value: Any, name: str, allow_empty: bool) -> set[int]: + if not isinstance(value, list): + raise TypeError(f"{name} must be a JSON array.") + if not value and not allow_empty: + raise ValueError(f"{name} must not be empty.") + block_ids = {_require_int(block_id, f"{name}[]", 0) for block_id in value} + if len(block_ids) != len(value): + raise ValueError(f"{name} must not contain duplicate block IDs.") + return block_ids + + +def _selector_from_trace_payload( + payload: dict[str, Any], + path: Path, +) -> tuple[SelectorMetrics, dict[str, Any]]: + if payload["schema_version"] != 1: + raise ValueError(f"{path}: selector trace schema_version must be 1.") + if payload["kind"] != "infllmv2_selector_samples": + raise ValueError(f"{path}: selector trace kind must be 'infllmv2_selector_samples', got {payload['kind']!r}.") + reference_definition = payload["reference_definition"] + if reference_definition != SELECTOR_REFERENCE_DEFINITION: + raise ValueError( + f"{path}: reference_definition must be {SELECTOR_REFERENCE_DEFINITION!r}, got {reference_definition!r}." + ) + rank = _require_int(payload["rank"], f"{path}: rank", 0) + layer = _require_int(payload["layer"], f"{path}: layer", 0) + block_size = _require_int(payload["block_size"], f"{path}: block_size", 1) + kernel_size = _require_int(payload["kernel_size"], f"{path}: kernel_size", 1) + kernel_stride = _require_int(payload["kernel_stride"], f"{path}: kernel_stride", 1) + topk = _require_int(payload["topk"], f"{path}: topk", 1) + dense_len = _require_int(payload["dense_len"], f"{path}: dense_len", 1) + init_blocks = _require_int(payload["init_blocks"], f"{path}: init_blocks", 0) + local_blocks = _require_int(payload["local_blocks"], f"{path}: local_blocks", 0) + selected_capacity = _require_int(payload["selected_capacity"], f"{path}: selected_capacity", 1) + if block_size % kernel_stride != 0: + raise ValueError(f"{path}: block_size must be divisible by kernel_stride.") + if block_size % (4 * kernel_stride) != 0: + raise ValueError(f"{path}: block_size must be divisible by 4 * kernel_stride.") + if dense_len < 4 * kernel_size: + raise ValueError(f"{path}: dense_len must be at least 4 * kernel_size.") + if init_blocks >= topk: + raise ValueError(f"{path}: init_blocks must be smaller than topk.") + expected_capacity = max( + topk + local_blocks, + (dense_len + block_size - 1) // block_size, + ) + if selected_capacity != expected_capacity: + raise ValueError( + f"{path}: selected_capacity must be {expected_capacity} for the " + f"header configuration, got {selected_capacity}." + ) + samples = payload["samples"] + if not isinstance(samples, list): + raise TypeError(f"{path}: samples must be a JSON array.") + if not samples: + raise ValueError(f"{path}: samples must not be empty.") + + hits = 0 + targets = 0 + selected_count_values: set[int] = set() + query_head_group_sizes: set[int] = set() + kv_heads: set[int] = set() + request_indices: set[int] = set() + query_positions: list[int] = [] + for sample_index, sample in enumerate(samples): + if not isinstance(sample, dict): + raise TypeError(f"{path}: samples[{sample_index}] must be a JSON object.") + sample_name = f"{path}: samples[{sample_index}]" + sample_rank = _require_int(sample["rank"], f"{sample_name}.rank", 0) + sample_layer = _require_int(sample["layer"], f"{sample_name}.layer", 0) + if sample_rank != rank or sample_layer != layer: + raise ValueError( + f"{sample_name}: rank/layer must match the trace header " + f"({rank}/{layer}), got {sample_rank}/{sample_layer}." + ) + request_index = _require_int(sample["request_index"], f"{sample_name}.request_index", 0) + _require_int(sample["query_index"], f"{sample_name}.query_index", 0) + _require_int(sample["query_offset"], f"{sample_name}.query_offset", 0) + query_position = _require_int(sample["query_position"], f"{sample_name}.query_position", 0) + kv_head = _require_int(sample["kv_head"], f"{sample_name}.kv_head", 0) + query_head_start = _require_int(sample["query_head_start"], f"{sample_name}.query_head_start", 0) + query_head_end = _require_int(sample["query_head_end"], f"{sample_name}.query_head_end", 1) + if query_head_end <= query_head_start: + raise ValueError(f"{sample_name}.query_head_end must be greater than query_head_start.") + query_head_group_size = query_head_end - query_head_start + if query_head_start != kv_head * query_head_group_size: + raise ValueError(f"{sample_name}: query-head range does not match contiguous GQA group {kv_head}.") + sample_block_size = _require_int(sample["block_size"], f"{sample_name}.block_size", 1) + sample_topk = _require_int(sample["topk"], f"{sample_name}.topk", 1) + if sample_block_size != block_size or sample_topk != topk: + raise ValueError( + f"{sample_name}: block_size/topk must match the trace header " + f"({block_size}/{topk}), got {sample_block_size}/{sample_topk}." + ) + selected_count = _require_int(sample["selected_count"], f"{sample_name}.selected_count", 1) + if selected_count > selected_capacity: + raise ValueError(f"{sample_name}.selected_count must not exceed selected_capacity {selected_capacity}.") + selected = _validate_block_ids( + sample["selected_blocks"], + f"{sample_name}.selected_blocks", + allow_empty=False, + ) + reference = _validate_block_ids( + sample["reference_blocks"], + f"{sample_name}.reference_blocks", + allow_empty=False, + ) + if len(selected) != selected_count or len(reference) != selected_count: + raise ValueError( + f"{sample_name}: selected_blocks and reference_blocks must both " + f"contain selected_count={selected_count} entries." + ) + if sample["selected_blocks"] != sorted(selected): + raise ValueError(f"{sample_name}.selected_blocks must be strictly increasing.") + if sample["reference_blocks"] != sorted(reference): + raise ValueError(f"{sample_name}.reference_blocks must be strictly increasing.") + if sample["reference_metric"] != reference_definition: + raise ValueError(f"{sample_name}.reference_metric must match reference_definition.") + valid_blocks = (query_position + 1 + block_size - 1) // block_size + if query_position + 1 < dense_len: + raise ValueError( + f"{sample_name}.query_position must be in the sparse region starting at visible length {dense_len}." + ) + if selected_count > valid_blocks: + raise ValueError(f"{sample_name}.selected_count must not exceed {valid_blocks} visible blocks.") + if max(selected.union(reference)) >= valid_blocks: + raise ValueError( + f"{sample_name}: block IDs must be smaller than the {valid_blocks} blocks visible at query_position." + ) + hits += len(selected.intersection(reference)) + targets += len(reference) + selected_count_values.add(selected_count) + query_head_group_sizes.add(query_head_group_size) + kv_heads.add(kv_head) + request_indices.add(request_index) + query_positions.append(query_position) + + if len(query_head_group_sizes) != 1: + raise ValueError(f"{path}: query-head group size must be consistent across selector samples.") + + selector = SelectorMetrics( + selector_hit_rate=hits / targets, + selector_samples=len(samples), + selector_hits=hits, + selector_targets=targets, + ) + config = { + "block_size": block_size, + "kernel_size": kernel_size, + "kernel_stride": kernel_stride, + "topk": topk, + "dense_len": dense_len, + "init_blocks": init_blocks, + "local_blocks": local_blocks, + "selected_capacity": selected_capacity, + "query_head_group_size": next(iter(query_head_group_sizes)), + "selected_count_values": sorted(selected_count_values), + "kv_heads_observed": sorted(kv_heads), + "request_slots_observed": sorted(request_indices), + "query_position_min": min(query_positions), + "query_position_max": max(query_positions), + } + provenance = { + "rank": rank, + "layer": layer, + "config": config, + "reference_definition": reference_definition, + } + return selector, provenance + + +def _load_selector_trace( + path: Path, +) -> tuple[SelectorMetrics, dict[str, Any], str, os.stat_result]: + payload, trace_sha256, trace_stat = _read_json_snapshot(path) + selector, provenance = _selector_from_trace_payload(payload, path) + return selector, provenance, trace_sha256, trace_stat + + +def _validate_workload_args(args: argparse.Namespace) -> None: + parsed_url = urlparse(args.base_url) + if parsed_url.scheme not in {"http", "https"} or not parsed_url.netloc: + raise ValueError(f"--base-url must be an absolute HTTP(S) URL, got {args.base_url!r}.") + if not args.endpoint.startswith("/"): + raise ValueError(f"--endpoint must start with '/', got {args.endpoint!r}.") + if not args.model.strip(): + raise ValueError("--model must not be empty.") + if args.tokenizer is not None and not args.tokenizer.strip(): + raise ValueError("--tokenizer must not be empty when provided.") + _require_int(args.seed, "--seed", 0) + _require_int(args.num_prompts, "--num-prompts", 1) + _require_int(args.input_len, "--input-len", 1) + _require_int(args.output_len, "--output-len", 2) + _require_int(args.max_concurrency, "--max-concurrency", 1) + if math.isnan(args.request_rate) or args.request_rate <= 0: + raise ValueError(f"--request-rate must be positive, got {args.request_rate}.") + + +def _validate_output_args(args: argparse.Namespace) -> None: + if args.output.exists() and not args.overwrite: + raise FileExistsError(f"Output already exists: {args.output}; pass --overwrite to replace it.") + if not args.output.parent.is_dir(): + raise FileNotFoundError(f"Output directory does not exist: {args.output.parent}") + + +def _validate_run_args(args: argparse.Namespace) -> None: + _validate_workload_args(args) + _validate_output_args(args) + _require_int(args.gpu_index, "--gpu-index", 0) + _require_number(args.sample_interval, "--sample-interval", 0.0, inclusive=False) + if args.variant == "dense" and args.selector_diagnostic is not None: + raise ValueError("--selector-diagnostic applies only to --variant sparse.") + if args.variant == "sparse" and args.selector_diagnostic is None: + raise ValueError("--variant sparse requires --selector-diagnostic.") + if args.variant == "sparse" and os.getenv(SELECTOR_TRACE_PATH_ENV): + raise ValueError( + f"Timed sparse runs require tracing to be disabled; unset {SELECTOR_TRACE_PATH_ENV} " + "and restart the sparse server without selector tracing." + ) + + +def _validate_diagnostic_args(args: argparse.Namespace) -> None: + _validate_workload_args(args) + _validate_output_args(args) + trace_path = args.trace_path.resolve() + if not trace_path.parent.is_dir(): + raise FileNotFoundError(f"Selector trace output directory does not exist: {trace_path.parent}") + if trace_path.exists(): + raise FileExistsError( + f"Selector trace output already exists: {trace_path}; diagnostic runs require a new trace path." + ) + configured_trace_path = os.getenv(SELECTOR_TRACE_PATH_ENV) + if configured_trace_path and Path(configured_trace_path).expanduser().resolve() != trace_path: + raise ValueError(f"--trace-path must match {SELECTOR_TRACE_PATH_ENV} when that environment variable is set.") + + +def _canonical_request_rate(request_rate: float) -> float | str: + return "inf" if math.isinf(request_rate) else request_rate + + +def _build_workload( + args: argparse.Namespace, + prompt_token_ids_sha256: str, + benchmark_serving_sha256: str, +) -> dict[str, Any]: + return { + "backend": "openai-chat", + "benchmark_serving_sha256": benchmark_serving_sha256, + "dataset_name": "random_token_ids", + "endpoint": args.endpoint, + "ignore_eos": True, + "input_len": args.input_len, + "max_concurrency": args.max_concurrency, + "model": args.model, + "num_prompts": args.num_prompts, + "output_len": args.output_len, + "prompt_token_ids_sha256": prompt_token_ids_sha256, + "random_range_ratio": 0.0, + "request_rate": _canonical_request_rate(args.request_rate), + "seed": args.seed, + "stream": True, + "tokenizer": args.tokenizer, + } + + +def _workload_id(workload: dict[str, Any]) -> str: + serialized = json.dumps(workload, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(serialized.encode("utf-8")).hexdigest() + + +def _prompt_token_ids_sha256(result: dict[str, Any], args: argparse.Namespace) -> str: + prompts = result["input_texts"] + if not isinstance(prompts, list): + raise TypeError("benchmark result: input_texts must be a JSON array.") + if len(prompts) != args.num_prompts: + raise RuntimeError(f"Expected {args.num_prompts} prompt token arrays, got {len(prompts)}.") + for prompt_index, prompt_token_ids in enumerate(prompts): + if not isinstance(prompt_token_ids, list): + raise TypeError(f"benchmark result: input_texts[{prompt_index}] must be a JSON array.") + if len(prompt_token_ids) != args.input_len: + raise RuntimeError( + f"Expected input_texts[{prompt_index}] to contain {args.input_len} " + f"tokens, got {len(prompt_token_ids)}." + ) + for token_index, token_id in enumerate(prompt_token_ids): + _require_int( + token_id, + f"benchmark result: input_texts[{prompt_index}][{token_index}]", + 0, + ) + serialized = json.dumps(prompts, ensure_ascii=False, separators=(",", ":")) + return hashlib.sha256(serialized.encode("utf-8")).hexdigest() + + +def _query_gpu_identity(gpu_index: int) -> GPUIdentity: + command = [ + "nvidia-smi", + f"--id={gpu_index}", + "--query-gpu=index,uuid,name,memory.total", + "--format=csv,noheader,nounits", + ] + completed = subprocess.run(command, check=True, capture_output=True, text=True) + rows = list(csv.reader(completed.stdout.splitlines(), skipinitialspace=True)) + if len(rows) != 1 or len(rows[0]) != 4: + raise RuntimeError(f"Expected one four-field nvidia-smi identity row for GPU {gpu_index}, got {rows!r}.") + reported_index = _require_int(int(rows[0][0]), f"GPU {gpu_index} reported index", 0) + if reported_index != gpu_index: + raise RuntimeError(f"nvidia-smi --id={gpu_index} reported physical GPU index {reported_index}.") + gpu_uuid = rows[0][1].strip() + gpu_name = rows[0][2].strip() + if not gpu_uuid or not gpu_name: + raise RuntimeError(f"GPU {gpu_index} UUID and name must not be empty.") + total_mib = _require_int(int(rows[0][3]), f"GPU {gpu_index} memory.total", 1) + return GPUIdentity( + gpu_index=reported_index, + gpu_uuid=gpu_uuid, + gpu_name=gpu_name, + total_mib=total_mib, + ) + + +def _query_gpu_memory_mib(gpu_index: int) -> int: + command = [ + "nvidia-smi", + f"--id={gpu_index}", + "--query-gpu=memory.used", + "--format=csv,noheader,nounits", + ] + completed = subprocess.run(command, check=True, capture_output=True, text=True) + lines = [line.strip() for line in completed.stdout.splitlines() if line.strip()] + if len(lines) != 1: + raise RuntimeError(f"Expected one nvidia-smi memory value for GPU {gpu_index}, got {lines!r}.") + return _require_int(int(lines[0]), f"GPU {gpu_index} memory.used", 0) + + +def _benchmark_command( + args: argparse.Namespace, + raw_result: Path, + *, + no_warmup: bool, +) -> list[str]: + command = [ + sys.executable, + str(BENCHMARK_SERVING), + "--backend", + "openai-chat", + "--base-url", + args.base_url.rstrip("/"), + "--endpoint", + args.endpoint, + "--model", + args.model, + "--dataset-name", + "random_token_ids", + "--seed", + str(args.seed), + "--num-prompts", + str(args.num_prompts), + "--random-input-len", + str(args.input_len), + "--random-output-len", + str(args.output_len), + "--random-range-ratio", + "0", + "--request-rate", + str(args.request_rate), + "--max-concurrency", + str(args.max_concurrency), + "--ignore-eos", + "--disable-tqdm", + "--percentile-metrics", + "ttft,tpot,itl", + "--metric-percentiles", + "99", + "--save-result", + "--result-filename", + str(raw_result), + ] + if args.tokenizer is not None: + command.extend(["--tokenizer", args.tokenizer]) + if no_warmup: + command.append("--no-warmup") + return command + + +def _run_with_memory_sampling(command: list[str], gpu_index: int, interval: float) -> GPUMemoryMetrics: + gpu_identity = _query_gpu_identity(gpu_index) + baseline_mib = _query_gpu_memory_mib(gpu_index) + memory_samples = [baseline_mib] + benchmark_process = subprocess.Popen(command) + try: + while benchmark_process.poll() is None: + memory_samples.append(_query_gpu_memory_mib(gpu_index)) + time.sleep(interval) + memory_samples.append(_query_gpu_memory_mib(gpu_index)) + if benchmark_process.returncode != 0: + raise subprocess.CalledProcessError(benchmark_process.returncode, command) + finally: + if benchmark_process.poll() is None: + benchmark_process.terminate() + try: + benchmark_process.wait(timeout=5) + except subprocess.TimeoutExpired: + benchmark_process.kill() + benchmark_process.wait() + + final_gpu_identity = _query_gpu_identity(gpu_index) + if final_gpu_identity != gpu_identity: + raise RuntimeError(f"GPU identity changed during benchmark: {gpu_identity} -> {final_gpu_identity}.") + peak_mib = max(memory_samples) + if peak_mib > gpu_identity.total_mib: + raise RuntimeError( + f"GPU {gpu_index} peak memory {peak_mib} MiB exceeds total memory {gpu_identity.total_mib} MiB." + ) + return GPUMemoryMetrics( + gpu_index=gpu_index, + gpu_uuid=gpu_identity.gpu_uuid, + gpu_name=gpu_identity.gpu_name, + total_mib=gpu_identity.total_mib, + baseline_mib=baseline_mib, + peak_mib=peak_mib, + peak_delta_mib=peak_mib - baseline_mib, + samples=len(memory_samples), + sample_interval_seconds=interval, + ) + + +def _validate_serving_result(result: dict[str, Any], args: argparse.Namespace) -> None: + completed = _require_int(result["completed"], "benchmark result: completed", 0) + if completed != args.num_prompts: + raise RuntimeError(f"Expected {args.num_prompts} completed requests, got {completed}.") + expected_input_tokens = args.num_prompts * args.input_len + actual_input_tokens = _require_int(result["total_input_tokens"], "benchmark result: total_input_tokens", 0) + if actual_input_tokens != expected_input_tokens: + raise RuntimeError(f"Expected {expected_input_tokens} input tokens, got {actual_input_tokens}.") + expected_output_tokens = args.num_prompts * args.output_len + actual_output_tokens = _require_int(result["total_output_tokens"], "benchmark result: total_output_tokens", 0) + if actual_output_tokens != expected_output_tokens: + raise RuntimeError(f"Expected {expected_output_tokens} output tokens, got {actual_output_tokens}.") + + for metric_name in ( + "duration", + "request_throughput", + "output_throughput", + "total_token_throughput", + "mean_ttft_ms", + "median_ttft_ms", + "p99_ttft_ms", + "mean_tpot_ms", + "median_tpot_ms", + "p99_tpot_ms", + ): + _require_number(result[metric_name], f"benchmark result: {metric_name}", 0.0, inclusive=False) + + +def _compact_serving_result(result: dict[str, Any]) -> dict[str, Any]: + return {field: result[field] for field in SERVING_RESULT_FIELDS} + + +def _verified_workload( + args: argparse.Namespace, + serving_result: dict[str, Any], + benchmark_serving_sha256_before: str, +) -> tuple[dict[str, Any], str]: + benchmark_serving_sha256_after = _sha256_file(BENCHMARK_SERVING) + if benchmark_serving_sha256_after != benchmark_serving_sha256_before: + raise RuntimeError(f"{BENCHMARK_SERVING} changed during the benchmark run.") + prompt_token_ids_sha256 = _prompt_token_ids_sha256(serving_result, args) + workload = _build_workload( + args, + prompt_token_ids_sha256=prompt_token_ids_sha256, + benchmark_serving_sha256=benchmark_serving_sha256_before, + ) + return workload, _workload_id(workload) + + +def _validate_trace_config(config: Any, name: str) -> dict[str, Any]: + if not isinstance(config, dict): + raise TypeError(f"{name} must be a JSON object.") + block_size = _require_int(config["block_size"], f"{name}.block_size", 1) + kernel_size = _require_int(config["kernel_size"], f"{name}.kernel_size", 1) + kernel_stride = _require_int(config["kernel_stride"], f"{name}.kernel_stride", 1) + topk = _require_int(config["topk"], f"{name}.topk", 1) + dense_len = _require_int(config["dense_len"], f"{name}.dense_len", 1) + init_blocks = _require_int(config["init_blocks"], f"{name}.init_blocks", 0) + local_blocks = _require_int(config["local_blocks"], f"{name}.local_blocks", 0) + selected_capacity = _require_int(config["selected_capacity"], f"{name}.selected_capacity", 1) + if block_size % kernel_stride != 0: + raise ValueError(f"{name}.block_size must be divisible by kernel_stride.") + if block_size % (4 * kernel_stride) != 0: + raise ValueError(f"{name}.block_size must be divisible by 4 * kernel_stride.") + if dense_len < 4 * kernel_size: + raise ValueError(f"{name}.dense_len must be at least 4 * kernel_size.") + if init_blocks >= topk: + raise ValueError(f"{name}.init_blocks must be smaller than topk.") + expected_capacity = max( + topk + local_blocks, + (dense_len + block_size - 1) // block_size, + ) + if selected_capacity != expected_capacity: + raise ValueError(f"{name}.selected_capacity must be {expected_capacity}, got {selected_capacity}.") + query_head_group_size = _require_int(config["query_head_group_size"], f"{name}.query_head_group_size", 1) + selected_count_values = config["selected_count_values"] + kv_heads_observed = config["kv_heads_observed"] + request_slots_observed = config["request_slots_observed"] + for values, field_name, minimum in ( + (selected_count_values, "selected_count_values", 1), + (kv_heads_observed, "kv_heads_observed", 0), + (request_slots_observed, "request_slots_observed", 0), + ): + if not isinstance(values, list) or not values: + raise ValueError(f"{name}.{field_name} must be a non-empty JSON array.") + normalized_values = [_require_int(value, f"{name}.{field_name}[]", minimum) for value in values] + if values != sorted(set(normalized_values)): + raise ValueError(f"{name}.{field_name} must be strictly increasing and unique.") + query_position_min = _require_int(config["query_position_min"], f"{name}.query_position_min", 0) + query_position_max = _require_int(config["query_position_max"], f"{name}.query_position_max", 0) + if query_position_min > query_position_max: + raise ValueError(f"{name}.query_position_min must not exceed query_position_max.") + if query_position_min + 1 < dense_len: + raise ValueError(f"{name}.query_position_min must lie in the sparse region.") + if any(value > selected_capacity for value in selected_count_values): + raise ValueError(f"{name}.selected_count_values must not exceed selected_capacity.") + return { + "block_size": block_size, + "kernel_size": kernel_size, + "kernel_stride": kernel_stride, + "topk": topk, + "dense_len": dense_len, + "init_blocks": init_blocks, + "local_blocks": local_blocks, + "selected_capacity": selected_capacity, + "query_head_group_size": query_head_group_size, + "selected_count_values": selected_count_values, + "kv_heads_observed": kv_heads_observed, + "request_slots_observed": request_slots_observed, + "query_position_min": query_position_min, + "query_position_max": query_position_max, + } + + +def _load_selector_diagnostic( + path: Path, +) -> tuple[dict[str, Any], str]: + diagnostic, diagnostic_sha256, _ = _read_json_snapshot(path) + if diagnostic["schema_version"] != SCHEMA_VERSION: + raise ValueError(f"{path}: selector diagnostic schema_version must be {SCHEMA_VERSION}.") + if diagnostic["kind"] != "infllmv2_selector_diagnostic": + raise ValueError(f"{path}: kind must be 'infllmv2_selector_diagnostic', got {diagnostic['kind']!r}.") + workload = diagnostic["workload"] + if not isinstance(workload, dict): + raise TypeError(f"{path}: workload must be a JSON object.") + _require_sha256( + workload["prompt_token_ids_sha256"], + f"{path}: workload.prompt_token_ids_sha256", + ) + _require_sha256( + workload["benchmark_serving_sha256"], + f"{path}: workload.benchmark_serving_sha256", + ) + workload_id = _workload_id(workload) + if diagnostic["workload_id"] != workload_id: + raise ValueError(f"{path}: workload_id does not match workload contents.") + + trace = diagnostic["trace"] + if not isinstance(trace, dict): + raise TypeError(f"{path}: trace must be a JSON object.") + _require_sha256(trace["sha256"], f"{path}: trace.sha256") + _require_int(trace["rank"], f"{path}: trace.rank", 0) + _require_int(trace["layer"], f"{path}: trace.layer", 0) + _validate_trace_config(trace["config"], f"{path}: trace.config") + if trace["reference_definition"] != SELECTOR_REFERENCE_DEFINITION: + raise ValueError(f"{path}: trace.reference_definition must be {SELECTOR_REFERENCE_DEFINITION!r}.") + + trace_window_start_ns = _require_int(diagnostic["trace_window_start_ns"], f"{path}: trace_window_start_ns", 0) + trace_window_end_ns = _require_int(diagnostic["trace_window_end_ns"], f"{path}: trace_window_end_ns", 0) + if trace_window_start_ns > trace_window_end_ns: + raise ValueError(f"{path}: trace request window is inverted.") + trace_mtime_ns = _require_int(trace["mtime_ns"], f"{path}: trace.mtime_ns", 0) + if not trace_window_start_ns <= trace_mtime_ns <= trace_window_end_ns: + raise ValueError(f"{path}: trace mtime falls outside the diagnostic request window.") + trace_path = Path(trace["path"]) + source_selector, source_provenance, source_trace_sha256, source_trace_stat = _load_selector_trace(trace_path) + if source_trace_sha256 != trace["sha256"]: + raise ValueError(f"{path}: trace SHA256 no longer matches {trace_path}.") + if source_trace_stat.st_mtime_ns != trace_mtime_ns: + raise ValueError(f"{path}: trace mtime no longer matches {trace_path}.") + if source_provenance != { + "rank": trace["rank"], + "layer": trace["layer"], + "config": trace["config"], + "reference_definition": trace["reference_definition"], + }: + raise ValueError(f"{path}: embedded trace provenance does not match {trace_path}.") + + selector = _validate_attached_selector(diagnostic["selector"]) + if selector != asdict(source_selector): + raise ValueError(f"{path}: embedded selector metrics do not match {trace_path}.") + if selector["selector_hits"] is None or selector["selector_targets"] is None: + raise ValueError(f"{path}: diagnostic selector must retain hit and target counts.") + if trace["sample_count"] != selector["selector_samples"]: + raise ValueError(f"{path}: trace.sample_count must match selector.selector_samples.") + if trace["metric_definition"] != SELECTOR_DEFINITION: + raise ValueError(f"{path}: trace.metric_definition must be {SELECTOR_DEFINITION!r}.") + return diagnostic, diagnostic_sha256 + + +def run_benchmark(args: argparse.Namespace) -> None: + _validate_run_args(args) + selector_diagnostic = None + selector_diagnostic_sha256 = None + if args.variant == "sparse": + selector_diagnostic, selector_diagnostic_sha256 = _load_selector_diagnostic(args.selector_diagnostic) + + benchmark_serving_sha256 = _sha256_file(BENCHMARK_SERVING) + with tempfile.TemporaryDirectory(prefix="infllmv2-benchmark-", dir=args.output.parent) as temporary_dir: + raw_result = Path(temporary_dir) / "serving.json" + command = _benchmark_command(args, raw_result, no_warmup=False) + print(f"Running with interpreter: {sys.executable}") + gpu_memory = _run_with_memory_sampling(command, args.gpu_index, args.sample_interval) + serving_result = _read_json(raw_result) + + _validate_serving_result(serving_result, args) + workload, workload_id = _verified_workload( + args, + serving_result, + benchmark_serving_sha256, + ) + print(f"Workload ID: {workload_id}") + + selector = None + selector_source = None + if selector_diagnostic is not None: + if selector_diagnostic["workload"] != workload: + raise ValueError("Selector diagnostic and sparse timing run do not describe the same workload.") + if selector_diagnostic["workload_id"] != workload_id: + raise ValueError("Selector diagnostic and sparse timing workload IDs differ.") + selector = selector_diagnostic["selector"] + trace = selector_diagnostic["trace"] + selector_source = { + "path": str(args.selector_diagnostic.resolve()), + "sha256": selector_diagnostic_sha256, + "workload_id": selector_diagnostic["workload_id"], + "trace_sha256": trace["sha256"], + "rank": trace["rank"], + "layer": trace["layer"], + "config": trace["config"], + "reference_definition": trace["reference_definition"], + } + + output = { + "schema_version": SCHEMA_VERSION, + "kind": "infllmv2_serving_run", + "variant": args.variant, + "workload_id": workload_id, + "workload": workload, + "gpu_memory": asdict(gpu_memory), + "selector": selector, + "selector_source": selector_source, + "serving": _compact_serving_result(serving_result), + } + _write_json(args.output, output, args.overwrite) + print(f"Saved {args.variant} result to {args.output}") + + +def run_selector_diagnostic(args: argparse.Namespace) -> None: + _validate_diagnostic_args(args) + trace_path = args.trace_path.resolve() + if trace_path.exists(): + raise FileExistsError(f"Selector trace output already exists immediately before the run: {trace_path}") + + benchmark_serving_sha256 = _sha256_file(BENCHMARK_SERVING) + trace_window_start_ns = time.time_ns() + with tempfile.TemporaryDirectory(prefix="infllmv2-diagnostic-", dir=args.output.parent) as temporary_dir: + raw_result = Path(temporary_dir) / "serving.json" + command = _benchmark_command(args, raw_result, no_warmup=True) + print(f"Running selector diagnostic with interpreter: {sys.executable}") + subprocess.run(command, check=True) + trace_window_end_ns = time.time_ns() + serving_result = _read_json(raw_result) + + _validate_serving_result(serving_result, args) + workload, workload_id = _verified_workload( + args, + serving_result, + benchmark_serving_sha256, + ) + selector, trace_provenance, trace_sha256, trace_stat = _load_selector_trace(trace_path) + if trace_stat.st_mtime_ns < trace_window_start_ns: + raise RuntimeError(f"Selector trace predates the diagnostic request window: {trace_path}") + if trace_stat.st_mtime_ns > trace_window_end_ns: + raise RuntimeError(f"Selector trace was modified after the diagnostic request window: {trace_path}") + + trace = { + "path": str(trace_path), + "sha256": trace_sha256, + "mtime_ns": trace_stat.st_mtime_ns, + "sample_count": selector.selector_samples, + "metric_definition": SELECTOR_DEFINITION, + **trace_provenance, + } + output = { + "schema_version": SCHEMA_VERSION, + "kind": "infllmv2_selector_diagnostic", + "workload_id": workload_id, + "workload": workload, + "trace_window_start_ns": trace_window_start_ns, + "trace_window_end_ns": trace_window_end_ns, + "trace": trace, + "selector": asdict(selector), + "serving": _compact_serving_result(serving_result), + } + _write_json(args.output, output, args.overwrite) + print(f"Workload ID: {workload_id}") + print(f"Selector block hit rate: {selector.selector_hit_rate:.2%} over {selector.selector_samples} samples") + print(f"Saved selector diagnostic to {args.output}") + + +def _load_run(path: Path, expected_variant: str) -> dict[str, Any]: + result = _read_json(path) + if result["schema_version"] != SCHEMA_VERSION: + raise ValueError(f"{path}: unsupported schema_version {result['schema_version']!r}.") + if result["kind"] != "infllmv2_serving_run": + raise ValueError(f"{path}: kind must be 'infllmv2_serving_run', got {result['kind']!r}.") + if result["variant"] != expected_variant: + raise ValueError(f"{path}: variant must be {expected_variant!r}, got {result['variant']!r}.") + workload = result["workload"] + if not isinstance(workload, dict): + raise TypeError(f"{path}: workload must be a JSON object.") + _require_sha256( + workload["prompt_token_ids_sha256"], + f"{path}: workload.prompt_token_ids_sha256", + ) + _require_sha256( + workload["benchmark_serving_sha256"], + f"{path}: workload.benchmark_serving_sha256", + ) + calculated_workload_id = _workload_id(workload) + if result["workload_id"] != calculated_workload_id: + raise ValueError(f"{path}: workload_id does not match the workload contents.") + if not isinstance(result["serving"], dict): + raise TypeError(f"{path}: serving must be a JSON object.") + if not isinstance(result["gpu_memory"], dict): + raise TypeError(f"{path}: gpu_memory must be a JSON object.") + if expected_variant == "dense": + if result["selector"] is not None or result["selector_source"] is not None: + raise ValueError(f"{path}: dense runs must not contain selector data.") + else: + _validate_attached_selector(result["selector"]) + selector_source = result["selector_source"] + if not isinstance(selector_source, dict): + raise TypeError(f"{path}: sparse selector_source must be a JSON object.") + if selector_source["workload_id"] != calculated_workload_id: + raise ValueError(f"{path}: selector_source workload_id does not match the run.") + _require_sha256(selector_source["sha256"], f"{path}: selector_source.sha256") + _require_sha256( + selector_source["trace_sha256"], + f"{path}: selector_source.trace_sha256", + ) + _require_int(selector_source["rank"], f"{path}: selector_source.rank", 0) + _require_int(selector_source["layer"], f"{path}: selector_source.layer", 0) + _validate_trace_config(selector_source["config"], f"{path}: selector_source.config") + if selector_source["reference_definition"] != SELECTOR_REFERENCE_DEFINITION: + raise ValueError(f"{path}: selector_source reference definition is invalid.") + selector_diagnostic_path = Path(selector_source["path"]) + selector_diagnostic, selector_diagnostic_sha256 = _load_selector_diagnostic(selector_diagnostic_path) + if selector_diagnostic_sha256 != selector_source["sha256"]: + raise ValueError(f"{path}: selector_source SHA256 no longer matches {selector_diagnostic_path}.") + if selector_diagnostic["workload"] != workload: + raise ValueError(f"{path}: selector diagnostic workload does not match the run.") + if selector_diagnostic["selector"] != result["selector"]: + raise ValueError(f"{path}: selector diagnostic metrics do not match the run.") + diagnostic_trace = selector_diagnostic["trace"] + if { + "trace_sha256": selector_source["trace_sha256"], + "rank": selector_source["rank"], + "layer": selector_source["layer"], + "config": selector_source["config"], + "reference_definition": selector_source["reference_definition"], + } != { + "trace_sha256": diagnostic_trace["sha256"], + "rank": diagnostic_trace["rank"], + "layer": diagnostic_trace["layer"], + "config": diagnostic_trace["config"], + "reference_definition": diagnostic_trace["reference_definition"], + }: + raise ValueError(f"{path}: selector_source provenance does not match the diagnostic.") + return result + + +def _validate_gpu_memory( + gpu_memory: dict[str, Any], + name: str, +) -> dict[str, Any]: + gpu_index = _require_int(gpu_memory["gpu_index"], f"{name}: gpu_index", 0) + gpu_uuid = gpu_memory["gpu_uuid"] + gpu_name = gpu_memory["gpu_name"] + if not isinstance(gpu_uuid, str) or not gpu_uuid: + raise ValueError(f"{name}: gpu_uuid must be a non-empty string.") + if not isinstance(gpu_name, str) or not gpu_name: + raise ValueError(f"{name}: gpu_name must be a non-empty string.") + total_mib = _require_int(gpu_memory["total_mib"], f"{name}: total_mib", 1) + baseline_mib = _require_int(gpu_memory["baseline_mib"], f"{name}: baseline_mib", 0) + peak_mib = _require_int(gpu_memory["peak_mib"], f"{name}: peak_mib", 1) + peak_delta_mib = _require_int(gpu_memory["peak_delta_mib"], f"{name}: peak_delta_mib", 0) + if baseline_mib > peak_mib: + raise ValueError(f"{name}: baseline_mib must not exceed peak_mib.") + if peak_mib > total_mib: + raise ValueError(f"{name}: peak_mib must not exceed total_mib.") + if peak_delta_mib != peak_mib - baseline_mib: + raise ValueError(f"{name}: peak_delta_mib must equal peak_mib - baseline_mib.") + _require_int(gpu_memory["samples"], f"{name}: samples", 1) + _require_number( + gpu_memory["sample_interval_seconds"], + f"{name}: sample_interval_seconds", + 0.0, + inclusive=False, + ) + return { + "gpu_index": gpu_index, + "gpu_uuid": gpu_uuid, + "gpu_name": gpu_name, + "total_mib": total_mib, + "baseline_mib": baseline_mib, + "peak_mib": peak_mib, + "peak_delta_mib": peak_delta_mib, + } + + +def _extract_report_metrics(run: dict[str, Any], name: str) -> dict[str, float | int]: + serving = run["serving"] + gpu_memory = _validate_gpu_memory(run["gpu_memory"], name) + return { + "completed": _require_int(serving["completed"], f"{name}: completed", 1), + "request_throughput_requests_per_second": _require_number( + serving["request_throughput"], f"{name}: request_throughput", 0.0, inclusive=False + ), + "output_throughput_tokens_per_second": _require_number( + serving["output_throughput"], f"{name}: output_throughput", 0.0, inclusive=False + ), + "total_throughput_tokens_per_second": _require_number( + serving["total_token_throughput"], f"{name}: total_token_throughput", 0.0, inclusive=False + ), + "mean_ttft_ms": _require_number(serving["mean_ttft_ms"], f"{name}: mean_ttft_ms", 0.0, inclusive=False), + "median_ttft_ms": _require_number(serving["median_ttft_ms"], f"{name}: median_ttft_ms", 0.0, inclusive=False), + "p99_ttft_ms": _require_number(serving["p99_ttft_ms"], f"{name}: p99_ttft_ms", 0.0, inclusive=False), + "mean_tpot_ms": _require_number(serving["mean_tpot_ms"], f"{name}: mean_tpot_ms", 0.0, inclusive=False), + "median_tpot_ms": _require_number(serving["median_tpot_ms"], f"{name}: median_tpot_ms", 0.0, inclusive=False), + "p99_tpot_ms": _require_number(serving["p99_tpot_ms"], f"{name}: p99_tpot_ms", 0.0, inclusive=False), + "baseline_gpu_memory_mib": gpu_memory["baseline_mib"], + "peak_gpu_memory_mib": gpu_memory["peak_mib"], + "peak_delta_gpu_memory_mib": gpu_memory["peak_delta_mib"], + } + + +def _ratio(sparse: float, dense: float) -> float: + return sparse / dense + + +def _reduction_percent(sparse: float, dense: float) -> float: + return (dense - sparse) / dense * 100.0 + + +def _validate_attached_selector(value: Any) -> dict[str, Any]: + if not isinstance(value, dict): + raise TypeError("Sparse run must contain selector metrics.") + hit_rate = _require_rate(value["selector_hit_rate"], "sparse selector_hit_rate") + sample_count = _require_int(value["selector_samples"], "sparse selector_samples", 1) + hits = value["selector_hits"] + targets = value["selector_targets"] + if (hits is None) != (targets is None): + raise ValueError("sparse selector_hits and selector_targets must either both be integers or both be null.") + if hits is not None: + hits = _require_int(hits, "sparse selector_hits", 0) + targets = _require_int(targets, "sparse selector_targets", 1) + if targets < sample_count: + raise ValueError( + "sparse selector_targets must be at least selector_samples because " + "every trace sample has a non-empty reference set." + ) + if hits > targets: + raise ValueError(f"sparse selector_hits ({hits}) must not exceed selector_targets ({targets}).") + calculated_hit_rate = hits / targets + if not math.isclose(hit_rate, calculated_hit_rate, rel_tol=0.0, abs_tol=1e-12): + raise ValueError( + f"sparse selector_hit_rate ({hit_rate}) does not match selector_hits / selector_targets " + f"({calculated_hit_rate})." + ) + definition = value["definition"] + if definition != SELECTOR_DEFINITION: + raise ValueError(f"Sparse selector definition must be {SELECTOR_DEFINITION!r}.") + return { + "selector_hit_rate": hit_rate, + "selector_samples": sample_count, + "selector_hits": hits, + "selector_targets": targets, + "definition": definition, + } + + +def _print_report(report: dict[str, Any]) -> None: + dense = report["metrics"]["dense"] + sparse = report["metrics"]["sparse"] + print(f"Workload ID: {report['workload_id']}") + print("| Metric | Dense | Sparse | Sparse / Dense |") + print("|---|---:|---:|---:|") + rows = ( + ( + "Request throughput (req/s)", + dense["request_throughput_requests_per_second"], + sparse["request_throughput_requests_per_second"], + ), + ( + "Output throughput (tok/s)", + dense["output_throughput_tokens_per_second"], + sparse["output_throughput_tokens_per_second"], + ), + ( + "Total token throughput (tok/s)", + dense["total_throughput_tokens_per_second"], + sparse["total_throughput_tokens_per_second"], + ), + ("Mean TTFT (ms)", dense["mean_ttft_ms"], sparse["mean_ttft_ms"]), + ("Median TTFT (ms)", dense["median_ttft_ms"], sparse["median_ttft_ms"]), + ("P99 TTFT (ms)", dense["p99_ttft_ms"], sparse["p99_ttft_ms"]), + ("Mean TPOT (ms)", dense["mean_tpot_ms"], sparse["mean_tpot_ms"]), + ("Median TPOT (ms)", dense["median_tpot_ms"], sparse["median_tpot_ms"]), + ("P99 TPOT (ms)", dense["p99_tpot_ms"], sparse["p99_tpot_ms"]), + ( + "Baseline GPU memory (MiB)", + dense["baseline_gpu_memory_mib"], + sparse["baseline_gpu_memory_mib"], + ), + ("Peak GPU memory (MiB)", dense["peak_gpu_memory_mib"], sparse["peak_gpu_memory_mib"]), + ( + "Peak delta GPU memory (MiB)", + dense["peak_delta_gpu_memory_mib"], + sparse["peak_delta_gpu_memory_mib"], + ), + ) + for label, dense_value, sparse_value in rows: + ratio = "N/A" if dense_value == 0 else f"{_ratio(sparse_value, dense_value):.3f}x" + print(f"| {label} | {dense_value:.3f} | {sparse_value:.3f} | {ratio} |") + selector = report["selector"] + print( + f"Selector block hit rate: {selector['selector_hit_rate']:.2%} " + f"over {selector['selector_samples']} request/query/KV-head samples" + ) + + +def create_report(args: argparse.Namespace) -> None: + if args.output.exists() and not args.overwrite: + raise FileExistsError(f"Output already exists: {args.output}; pass --overwrite to replace it.") + dense_run = _load_run(args.dense_result, "dense") + sparse_run = _load_run(args.sparse_result, "sparse") + if dense_run["workload"] != sparse_run["workload"]: + raise ValueError("Dense and sparse results do not describe the same workload.") + if dense_run["workload_id"] != sparse_run["workload_id"]: + raise ValueError("Dense and sparse workload IDs differ.") + + dense_gpu = _validate_gpu_memory(dense_run["gpu_memory"], "dense") + sparse_gpu = _validate_gpu_memory(sparse_run["gpu_memory"], "sparse") + if dense_gpu["gpu_uuid"] != sparse_gpu["gpu_uuid"]: + raise ValueError( + "Dense and sparse runs must use the same physical GPU UUID, got " + f"{dense_gpu['gpu_uuid']!r} and {sparse_gpu['gpu_uuid']!r}." + ) + if dense_gpu["gpu_name"] != sparse_gpu["gpu_name"] or dense_gpu["total_mib"] != sparse_gpu["total_mib"]: + raise ValueError("Dense and sparse GPU identity metadata differs.") + + dense_metrics = _extract_report_metrics(dense_run, "dense") + sparse_metrics = _extract_report_metrics(sparse_run, "sparse") + selector = _validate_attached_selector(sparse_run["selector"]) + if selector["definition"] != SELECTOR_DEFINITION: + raise ValueError(f"Sparse selector definition must be {SELECTOR_DEFINITION!r}.") + + comparison = { + "output_throughput_speedup": _ratio( + sparse_metrics["output_throughput_tokens_per_second"], + dense_metrics["output_throughput_tokens_per_second"], + ), + "request_throughput_speedup": _ratio( + sparse_metrics["request_throughput_requests_per_second"], + dense_metrics["request_throughput_requests_per_second"], + ), + "total_token_throughput_speedup": _ratio( + sparse_metrics["total_throughput_tokens_per_second"], + dense_metrics["total_throughput_tokens_per_second"], + ), + "mean_ttft_reduction_percent": _reduction_percent( + sparse_metrics["mean_ttft_ms"], dense_metrics["mean_ttft_ms"] + ), + "mean_tpot_reduction_percent": _reduction_percent( + sparse_metrics["mean_tpot_ms"], dense_metrics["mean_tpot_ms"] + ), + "peak_gpu_memory_reduction_percent": _reduction_percent( + sparse_metrics["peak_gpu_memory_mib"], dense_metrics["peak_gpu_memory_mib"] + ), + } + report = { + "schema_version": SCHEMA_VERSION, + "kind": "infllmv2_dense_sparse_report", + "workload_id": dense_run["workload_id"], + "workload": dense_run["workload"], + "gpu": { + "uuid": dense_gpu["gpu_uuid"], + "name": dense_gpu["gpu_name"], + "total_mib": dense_gpu["total_mib"], + "dense_index": dense_gpu["gpu_index"], + "sparse_index": sparse_gpu["gpu_index"], + }, + "metrics": {"dense": dense_metrics, "sparse": sparse_metrics}, + "comparison": comparison, + "selector": selector, + "selector_source": sparse_run["selector_source"], + } + _write_json(args.output, report, args.overwrite) + _print_report(report) + print(f"Saved comparison report to {args.output}") + + +def summarize_selector(args: argparse.Namespace) -> None: + selector, provenance, trace_sha256, _ = _load_selector_trace(args.samples) + output = { + "schema_version": SCHEMA_VERSION, + "kind": "infllmv2_unbound_selector_summary", + "trace_sha256": trace_sha256, + "trace": provenance, + **asdict(selector), + } + _write_json(args.output, output, args.overwrite) + print(f"Selector block hit rate: {selector.selector_hit_rate:.2%} over {selector.selector_samples} samples") + print(f"Saved selector summary to {args.output}") + + +def _parse_operator_scenario(value: str) -> OperatorScenario: + parts = value.split(":") + if len(parts) != 2: + raise argparse.ArgumentTypeError("operator scenario must be CONTEXT_LENGTH:CONCURRENCY, for example 32768:1") + try: + context_length = int(parts[0]) + concurrency = int(parts[1]) + except ValueError as exc: + raise argparse.ArgumentTypeError("operator scenario context length and concurrency must be integers") from exc + if context_length < 1: + raise argparse.ArgumentTypeError("operator scenario context length must be positive") + if concurrency < 1: + raise argparse.ArgumentTypeError("operator scenario concurrency must be positive") + return OperatorScenario(context_length=context_length, concurrency=concurrency) + + +def _percentile(sorted_values: list[float], quantile: float) -> float: + index = max(0, math.ceil(quantile * len(sorted_values)) - 1) + return sorted_values[index] + + +def _time_gpu_operator(paddle, operation, warmup: int, repeats: int) -> dict[str, Any]: + for _ in range(warmup): + operation() + paddle.device.synchronize() + starts = [paddle.device.cuda.Event(enable_timing=True) for _ in range(repeats)] + ends = [paddle.device.cuda.Event(enable_timing=True) for _ in range(repeats)] + wall_start = time.perf_counter() + for start, end in zip(starts, ends): + start.record() + operation() + end.record() + paddle.device.synchronize() + wall_mean_us = (time.perf_counter() - wall_start) * 1e6 / repeats + elapsed_us = sorted(start.elapsed_time(end) * 1000.0 for start, end in zip(starts, ends)) + return { + "mean_us": statistics.fmean(elapsed_us), + "median_us": statistics.median(elapsed_us), + "p10_us": _percentile(elapsed_us, 0.10), + "p90_us": _percentile(elapsed_us, 0.90), + "p99_us": _percentile(elapsed_us, 0.99), + "min_us": elapsed_us[0], + "max_us": elapsed_us[-1], + "wall_mean_us": wall_mean_us, + "warmup": warmup, + "repeats": repeats, + } + + +def _run_operator_scenario(paddle, sparse_ops, dense_ops, args, scenario: OperatorScenario) -> dict[str, Any]: + update_compressed_k, select_blocks, attention_forward = sparse_ops + append_attention, get_block_shape_and_split_kv_block = dense_ops + block_size = 64 + query_heads = 32 + kv_heads = 2 + head_dim = 128 + kernel_size = 32 + kernel_stride = 16 + topk = 64 + dense_len = 8192 + init_blocks = 1 + local_blocks = 32 + selected_capacity = max( + topk + local_blocks, + (dense_len + block_size - 1) // block_size, + ) + max_blocks_per_seq = (scenario.context_length + block_size - 1) // block_size + if max_blocks_per_seq > 2048: + raise ValueError("InfLLM-V2 operator scenarios support at most 2048 blocks (131072 tokens at block_size=64).") + physical_blocks = scenario.concurrency * max_blocks_per_seq + kv_splits = (selected_capacity + STAGE2_BLOCKS_PER_SPLIT - 1) // STAGE2_BLOCKS_PER_SPLIT + + paddle.seed(args.seed) + query = paddle.randn([scenario.concurrency, query_heads, head_dim], dtype="float32").astype(args.dtype) + key_cache = paddle.zeros([physical_blocks, kv_heads, block_size, head_dim], dtype=args.dtype) + value_cache = paddle.zeros_like(key_cache) + compressed_k = paddle.zeros( + [physical_blocks, kv_heads, block_size // kernel_stride, head_dim], + dtype=args.dtype, + ) + compressed_k2 = paddle.zeros( + [ + physical_blocks, + kv_heads, + block_size // (4 * kernel_stride), + head_dim, + ], + dtype=args.dtype, + ) + block_tables = paddle.arange(physical_blocks, dtype="int32").reshape([scenario.concurrency, max_blocks_per_seq]) + seq_lens_decoder = paddle.full([scenario.concurrency], scenario.context_length - 1, dtype="int32") + seq_lens_this_time = paddle.ones([scenario.concurrency], dtype="int32") + batch_id_per_token = paddle.arange(scenario.concurrency, dtype="int32") + cu_seqlens_q = paddle.arange(scenario.concurrency + 1, dtype="int32") + topk_indices = paddle.empty([scenario.concurrency, kv_heads, selected_capacity], dtype="int32") + block_scores = paddle.empty([scenario.concurrency, kv_heads, max_blocks_per_seq], dtype="float32") + selected_counts = paddle.empty([scenario.concurrency, kv_heads], dtype="int32") + coarse_lse = paddle.empty([scenario.concurrency, query_heads], dtype="float32") + max_coarse_windows = max( + 0, + (max_blocks_per_seq * block_size - 4 * kernel_size) // (4 * kernel_stride) + 1, + ) + coarse_splits = max(1, (max_coarse_windows + 15) // 16) + coarse_partial_max = paddle.empty([scenario.concurrency, query_heads, coarse_splits], dtype="float32") + coarse_partial_sum = paddle.empty([scenario.concurrency, query_heads, coarse_splits], dtype="float32") + attention_out = paddle.empty(query.shape, dtype=args.dtype) + partial_acc = paddle.empty( + [scenario.concurrency, query_heads, kv_splits, head_dim], + dtype="float32", + ) + partial_max = paddle.empty([scenario.concurrency, query_heads, kv_splits], dtype="float32") + partial_sum = paddle.empty([scenario.concurrency, query_heads, kv_splits], dtype="float32") + + if args.dtype == "float32": + raise ValueError("FastDeploy decode_unified_attention dense baseline does not support float32.") + group_size = query_heads // kv_heads + dense_max_lengths = paddle.zeros([6], dtype="int32").cpu() + seq_lens_encoder = paddle.zeros([scenario.concurrency], dtype="int32") + max_decoder_tiles = scenario.concurrency * (scenario.context_length * group_size + 15) // 16 + decoder_batch_ids = paddle.zeros([max_decoder_tiles], dtype="int32") + decoder_tile_ids = paddle.zeros_like(decoder_batch_ids) + decoder_num_blocks_cpu = paddle.zeros([1], dtype="int32").cpu() + decoder_num_blocks_device = paddle.zeros([1], dtype="int32") + decoder_chunk_size_device = paddle.zeros([1], dtype="int32") + max_encoder_tiles = scenario.concurrency * (scenario.context_length * group_size + 63) // 64 + encoder_batch_ids = paddle.zeros([max_encoder_tiles], dtype="int32") + encoder_tile_ids = paddle.zeros_like(encoder_batch_ids) + encoder_num_blocks_cpu = paddle.zeros([1], dtype="int32").cpu() + kv_batch_ids = paddle.zeros([max_encoder_tiles], dtype="int32") + kv_tile_ids = paddle.zeros_like(kv_batch_ids) + kv_num_blocks_cpu = paddle.zeros([1], dtype="int32").cpu() + get_block_shape_and_split_kv_block( + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + decoder_batch_ids, + decoder_tile_ids, + decoder_num_blocks_cpu, + decoder_num_blocks_device, + decoder_chunk_size_device, + dense_max_lengths, + encoder_batch_ids, + encoder_tile_ids, + encoder_num_blocks_cpu, + kv_batch_ids, + kv_tile_ids, + kv_num_blocks_cpu, + 64, + 16, + group_size, + block_size, + ) + current_kv = paddle.zeros([scenario.concurrency, 2 * kv_heads * head_dim], dtype=args.dtype) + dense_qkv = paddle.concat( + [query.reshape([scenario.concurrency, query_heads * head_dim]), current_kv], + axis=1, + ) + + def update_op(): + return update_compressed_k( + query, + key_cache, + compressed_k, + compressed_k2, + block_tables, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + kernel_size, + kernel_stride, + ) + + def stage1_op(): + return select_blocks( + query, + compressed_k, + compressed_k2, + block_tables, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + topk_indices, + block_scores, + selected_counts, + coarse_lse, + coarse_partial_max, + coarse_partial_sum, + block_size, + kernel_size, + kernel_stride, + topk, + dense_len, + init_blocks, + local_blocks, + ) + + def stage2_op(): + return attention_forward( + query, + key_cache, + value_cache, + block_tables, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + topk_indices, + attention_out, + partial_acc, + partial_max, + partial_sum, + ) + + def sparse_chain(): + update_op() + stage1_op() + stage2_op() + + def dense_attention_op(): + return append_attention( + dense_qkv, + key_cache, + value_cache, + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + block_tables, + encoder_batch_ids, + encoder_tile_ids, + encoder_num_blocks_cpu, + kv_batch_ids, + kv_tile_ids, + kv_num_blocks_cpu, + decoder_batch_ids, + decoder_tile_ids, + decoder_num_blocks_cpu, + dense_max_lengths, + rotary_embs=None, + attn_mask=None, + qkv_bias=None, + qkv_scale=None, + k_quant_scale=None, + v_quant_scale=None, + k_dequant_scale=None, + v_dequant_scale=None, + cache_k_zp=None, + cache_v_zp=None, + linear_shift=None, + linear_smooth=None, + mask_offset=None, + kv_signal_data=None, + q_norm_weight=None, + k_norm_weight=None, + sinks=None, + rms_norm_eps=1e-6, + compute_type="bf16" if args.dtype == "bfloat16" else "fp16", + cache_quant_type="none", + use_neox_rotary_style=False, + rope_3d=False, + max_input_length=scenario.context_length, + quant_max_bound=0.0, + quant_min_bound=0.0, + out_linear_in_scale=-1.0, + encoder_block_shape_q=64, + decoder_block_shape_q=16, + # Match FlashAttentionBackend and the serving runs. A single + # 32K partition severely underfills the SMs and is not an + # equivalent dense-decode baseline. + max_partition_size=1024, + encoder_max_partition_size=scenario.context_length, + speculate_max_draft_token_num=1, + causal=True, + speculate_decoder=False, + sliding_window=0, + sink_size=0, + head_wise_full_hidden=0, + only_do_attn=True, + ) + + update_op() + stage1_outputs = stage1_op() + stage2_outputs = stage2_op() + dense_attention_op() + paddle.device.synchronize() + if not isinstance(stage1_outputs, (tuple, list)) or len(stage1_outputs) != 6: + raise RuntimeError("infllmv2_select_blocks must return six in-place workspace aliases.") + expected_stage1_workspaces = ( + topk_indices, + block_scores, + selected_counts, + coarse_lse, + coarse_partial_max, + coarse_partial_sum, + ) + for output_index, (returned, workspace) in enumerate(zip(stage1_outputs, expected_stage1_workspaces)): + if not returned._is_shared_buffer_with(workspace): + raise RuntimeError( + f"infllmv2_select_blocks output {output_index} does not alias its persistent input workspace." + ) + if not isinstance(stage2_outputs, (tuple, list)) or len(stage2_outputs) != 4: + raise RuntimeError("infllmv2_attention_forward must return four in-place workspace aliases.") + expected_workspaces = (attention_out, partial_acc, partial_max, partial_sum) + for output_index, (returned, workspace) in enumerate(zip(stage2_outputs, expected_workspaces)): + if not returned._is_shared_buffer_with(workspace): + raise RuntimeError( + f"infllmv2_attention_forward output {output_index} does not alias its persistent input workspace." + ) + + count_values = selected_counts.numpy().reshape(-1).tolist() + if len(set(count_values)) != 1: + raise RuntimeError("Synthetic operator workload must select the same block count for every request/KV head.") + selected_blocks = int(count_values[0]) + selected_indices = topk_indices.numpy() + effective_stage2_tokens = 0 + for token_id in range(scenario.concurrency): + visible_tokens = scenario.context_length + for kv_head in range(kv_heads): + for logical_block in selected_indices[token_id, kv_head]: + if logical_block < 0: + break + block_start = int(logical_block) * block_size + effective_stage2_tokens += max(0, min(block_size, visible_tokens - block_start)) + visible_stage2_tokens = scenario.concurrency * kv_heads * scenario.context_length + effective_sparse_kv_fraction = effective_stage2_tokens / visible_stage2_tokens + preset_sparse_kv_fraction = min(1.0, (topk + local_blocks) * block_size / scenario.context_length) + selected_tokens = effective_stage2_tokens // (scenario.concurrency * kv_heads) + workspace_tensors = ( + topk_indices, + block_scores, + selected_counts, + coarse_lse, + coarse_partial_max, + coarse_partial_sum, + attention_out, + partial_acc, + partial_max, + partial_sum, + ) + stage2_workspace_tensors = ( + attention_out, + partial_acc, + partial_max, + partial_sum, + ) + workspace_bytes = sum(math.prod(tensor.shape) * tensor.element_size() for tensor in workspace_tensors) + stage2_workspace_bytes = sum( + math.prod(tensor.shape) * tensor.element_size() for tensor in stage2_workspace_tensors + ) + return { + "context_length": scenario.context_length, + "concurrency": scenario.concurrency, + "max_blocks_per_seq": max_blocks_per_seq, + "selected_blocks": selected_blocks, + "selected_tokens": selected_tokens, + "effective_stage2_kv_tokens": effective_stage2_tokens, + "visible_stage2_kv_tokens": visible_stage2_tokens, + "effective_sparse_kv_fraction": effective_sparse_kv_fraction, + "preset_sparse_kv_fraction": preset_sparse_kv_fraction, + "effective_minus_preset_sparse_kv_fraction": (effective_sparse_kv_fraction - preset_sparse_kv_fraction), + "configured_topk_only_kv_fraction": min(1.0, topk * block_size / scenario.context_length), + "stage2_token_sparsity_ideal_speedup": 1.0 / effective_sparse_kv_fraction, + "persistent_workspace_bytes": workspace_bytes, + "persistent_stage2_workspace_bytes": stage2_workspace_bytes, + "stage2_inplace_alias_verified": True, + "operators": { + "update_compressed_k": _time_gpu_operator(paddle, update_op, args.warmup, args.repeats), + "stage1_select_blocks": _time_gpu_operator(paddle, stage1_op, args.warmup, args.repeats), + "stage2_attention": _time_gpu_operator(paddle, stage2_op, args.warmup, args.repeats), + "sparse_decode_chain": _time_gpu_operator(paddle, sparse_chain, args.warmup, args.repeats), + "dense_attention": _time_gpu_operator(paddle, dense_attention_op, args.warmup, args.repeats), + }, + } + + +def run_operator_benchmark(args: argparse.Namespace) -> None: + _validate_output_args(args) + _require_int(args.seed, "--seed", 0) + _require_int(args.warmup, "--warmup", 1) + _require_int(args.repeats, "--repeats", 1) + _require_int(args.gpu_index, "--gpu-index", 0) + if args.scenarios is None: + args.scenarios = [ + OperatorScenario(context_length=32768, concurrency=1), + OperatorScenario(context_length=131072, concurrency=4), + ] + scenario_keys = {(scenario.context_length, scenario.concurrency) for scenario in args.scenarios} + if len(scenario_keys) != len(args.scenarios): + raise ValueError("--scenario entries must be unique.") + + import paddle + + if not paddle.is_compiled_with_cuda(): + raise RuntimeError("The InfLLM-V2 operator benchmark requires CUDA Paddle.") + paddle.set_device(args.device) + from fastdeploy.model_executor.layers.attention.ops import ( + append_attention, + get_block_shape_and_split_kv_block, + ) + from fastdeploy.model_executor.ops.gpu import ( + infllmv2_attention_forward, + infllmv2_select_blocks, + infllmv2_update_compressed_k, + ) + + gpu_identity = _query_gpu_identity(args.gpu_index) + scenarios = [] + for scenario in args.scenarios: + scenarios.append( + _run_operator_scenario( + paddle, + ( + infllmv2_update_compressed_k, + infllmv2_select_blocks, + infllmv2_attention_forward, + ), + (append_attention, get_block_shape_and_split_kv_block), + args, + scenario, + ) + ) + paddle.device.cuda.empty_cache() + + output = { + "schema_version": SCHEMA_VERSION, + "kind": "infllmv2_operator_benchmark", + "gpu": asdict(gpu_identity), + "device": args.device, + "dtype": args.dtype, + "seed": args.seed, + "model_shape": { + "query_heads": 32, + "kv_heads": 2, + "gqa_group_size": 16, + "head_dim": 128, + "block_size": 64, + "selected_capacity": 128, + "blocks_per_stage2_split": STAGE2_BLOCKS_PER_SPLIT, + }, + "source_sha256": { + "launcher": _sha256_file(Path(__file__).parents[1] / "custom_ops/gpu_ops/infllmv2_attention/infllmv2.cu"), + "kernels": _sha256_file( + Path(__file__).parents[1] / "custom_ops/gpu_ops/infllmv2_attention/infllmv2_impl.cuh" + ), + }, + "loaded_fastdeploy_ops": _loaded_shared_object_metadata("fastdeploy_ops_pd_.so"), + "scenarios": scenarios, + } + _write_json(args.output, output, args.overwrite) + for scenario in scenarios: + operators = scenario["operators"] + print( + f"{scenario['context_length']} tokens / concurrency " + f"{scenario['concurrency']}: " + f"update={operators['update_compressed_k']['median_us']:.3f} us, " + f"stage1={operators['stage1_select_blocks']['median_us']:.3f} us, " + f"stage2={operators['stage2_attention']['median_us']:.3f} us, " + f"chain={operators['sparse_decode_chain']['median_us']:.3f} us, " + f"dense={operators['dense_attention']['median_us']:.3f} us, " + f"effective_sparse_kv_fraction={scenario['effective_sparse_kv_fraction']:.6f}" + ) + print(f"Saved operator benchmark to {args.output}") + + +def _time_prefill_operator(paddle, operation, warmup: int, repeats: int) -> dict[str, float | int]: + for _ in range(warmup): + operation() + paddle.device.synchronize() + elapsed = [] + for _ in range(repeats): + start = paddle.device.cuda.Event(enable_timing=True) + end = paddle.device.cuda.Event(enable_timing=True) + start.record() + operation() + end.record() + end.synchronize() + elapsed.append(start.elapsed_time(end)) + elapsed.sort() + return { + "median_ms": statistics.median(elapsed), + "mean_ms": statistics.fmean(elapsed), + "min_ms": elapsed[0], + "max_ms": elapsed[-1], + "warmup": warmup, + "repeats": repeats, + } + + +def _make_prefill_backend(backend_type, fine, coarse): + backend = backend_type.__new__(backend_type) + backend.block_size = 64 + backend.num_heads = 32 + backend.kv_num_heads = 2 + backend.head_dim = 128 + backend.kernel_size = 32 + backend.kernel_stride = 16 + backend.topk = 64 + backend.local_blocks = 32 + backend.init_blocks = 1 + backend.dense_len = 8192 + backend.selected_capacity = 128 + backend.prefill_query_chunk_size = 4096 + backend._compressed_k = fine + backend._compressed_k2 = coarse + for name in ( + "_workspace_key", + "_topk_indices_ws", + "_block_scores_ws", + "_selected_counts_ws", + "_coarse_lse_ws", + "_coarse_partial_max_ws", + "_coarse_partial_sum_ws", + "_attention_out_ws", + "_partial_acc_ws", + "_partial_max_ws", + "_partial_sum_ws", + ): + setattr(backend, name, None) + return backend + + +def _run_prefill_context(paddle, context_length: int, dtype: str, warmup: int, repeats: int): + from fastdeploy.model_executor.layers.attention.flash_attn_backend import ( + flash_attn_func, + ) + from fastdeploy.model_executor.layers.attention.infllmv2_attention_backend import ( + InfLLMV2AttentionBackend, + ) + from fastdeploy.model_executor.ops.gpu import infllmv2_update_compressed_k + + block_size = 64 + query_heads = 32 + kv_heads = 2 + head_dim = 128 + dense_len = 8192 + if context_length <= dense_len: + raise ValueError(f"context length must exceed dense_len={dense_len}") + blocks = (context_length + block_size - 1) // block_size + padded_length = blocks * block_size + + query = paddle.randn([context_length, query_heads, head_dim], dtype=dtype) + key = paddle.randn([context_length, kv_heads, head_dim], dtype=dtype) + value = paddle.randn(key.shape, dtype=dtype) + padded_key = paddle.zeros([padded_length, kv_heads, head_dim], dtype=dtype) + padded_value = paddle.zeros_like(padded_key) + padded_key[:context_length] = key + padded_value[:context_length] = value + key_cache = paddle.transpose( + padded_key.reshape([blocks, block_size, kv_heads, head_dim]), [0, 2, 1, 3] + ).contiguous() + value_cache = paddle.transpose( + padded_value.reshape([blocks, block_size, kv_heads, head_dim]), [0, 2, 1, 3] + ).contiguous() + block_tables = paddle.arange(blocks, dtype="int32").reshape([1, blocks]) + fine = paddle.zeros([blocks, kv_heads, 4, head_dim], dtype=dtype) + coarse = paddle.zeros([blocks, kv_heads, 1, head_dim], dtype=dtype) + seq_lens_decoder = paddle.zeros([1], dtype="int32") + seq_lens_this_time = paddle.to_tensor([context_length], dtype="int32") + batch_ids = paddle.zeros([context_length], dtype="int32") + cu_seqlens = paddle.to_tensor([0, context_length], dtype="int32") + backend = _make_prefill_backend(InfLLMV2AttentionBackend, fine, coarse) + + def update_summaries(): + return infllmv2_update_compressed_k( + query, + key_cache, + fine, + coarse, + block_tables, + seq_lens_decoder, + seq_lens_this_time, + batch_ids, + cu_seqlens, + 32, + 16, + ) + + def sparse_without_update(): + return backend._sparse_prefill_attention( + query, + key, + value, + key_cache, + value_cache, + block_tables, + ) + + def sparse_chain(): + update_summaries() + return sparse_without_update() + + def dense_attention(): + return flash_attn_func( + query, + key, + value, + cu_seqlens, + cu_seqlens, + context_length, + context_length, + causal=True, + num_heads=query_heads, + kv_num_heads=kv_heads, + head_dim=head_dim, + version=2, + )[0] + + update_summaries() + sparse_without_update() + dense_attention() + paddle.device.synchronize() + operators = { + "update_compressed_k": _time_prefill_operator(paddle, update_summaries, warmup, repeats), + "sparse_without_update": _time_prefill_operator(paddle, sparse_without_update, warmup, repeats), + "sparse_prefill_chain": _time_prefill_operator(paddle, sparse_chain, warmup, repeats), + "dense_prefill": _time_prefill_operator(paddle, dense_attention, warmup, repeats), + } + dense_ms = operators["dense_prefill"]["median_ms"] + sparse_ms = operators["sparse_prefill_chain"]["median_ms"] + return { + "context_length": context_length, + "dense_len": dense_len, + "query_tile_size": 128, + "query_chunk_size": 4096, + "selected_blocks": 96, + "operators": operators, + "dense_over_sparse_speedup": dense_ms / sparse_ms, + } + + +def run_prefill_benchmark(args: argparse.Namespace) -> None: + _validate_output_args(args) + _require_int(args.seed, "--seed", 0) + _require_int(args.warmup, "--warmup", 1) + _require_int(args.repeats, "--repeats", 1) + context_lengths = args.context_lengths or [16384, 32768] + if len(set(context_lengths)) != len(context_lengths): + raise ValueError("--context-length entries must be unique") + + import paddle + + if not paddle.is_compiled_with_cuda(): + raise RuntimeError("CUDA Paddle is required") + paddle.set_device(args.device) + paddle.seed(args.seed) + scenarios = [ + _run_prefill_context(paddle, length, args.dtype, args.warmup, args.repeats) for length in context_lengths + ] + root = Path(__file__).parents[1] + properties = paddle.device.get_device_properties() + output = { + "kind": "infllmv2_sparse_prefill_operator_benchmark", + "device": args.device, + "gpu_name": properties.name, + "dtype": args.dtype, + "seed": args.seed, + "source_sha256": { + "backend": _sha256_file(root / "fastdeploy/model_executor/layers/attention/infllmv2_attention_backend.py"), + "benchmark": _sha256_file(Path(__file__)), + }, + "scenarios": scenarios, + } + _write_json(args.output, output, args.overwrite) + for scenario in scenarios: + print( + f"{scenario['context_length']} tokens: " + f"dense={scenario['operators']['dense_prefill']['median_ms']:.3f} ms, " + f"sparse_chain={scenario['operators']['sparse_prefill_chain']['median_ms']:.3f} ms, " + f"speedup={scenario['dense_over_sparse_speedup']:.3f}x" + ) + print(f"Saved benchmark to {args.output}") + + +CUDA_IMPL_BLOCK_SIZE = 64 +CUDA_IMPL_QUERY_HEADS = 32 +CUDA_IMPL_KV_HEADS = 2 +CUDA_IMPL_HEAD_DIM = 128 +CUDA_IMPL_KERNEL_SIZE = 32 +CUDA_IMPL_KERNEL_STRIDE = 16 +CUDA_IMPL_TOPK = 64 +CUDA_IMPL_LOCAL_BLOCKS = 32 +CUDA_IMPL_INIT_BLOCKS = 1 +CUDA_IMPL_SELECTED_BLOCKS = CUDA_IMPL_TOPK + CUDA_IMPL_LOCAL_BLOCKS + + +def _time_torch_cuda(torch, operation, warmup: int, repeats: int) -> dict[str, Any]: + for _ in range(warmup): + operation() + torch.cuda.synchronize() + starts = [torch.cuda.Event(enable_timing=True) for _ in range(repeats)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(repeats)] + wall_start = time.perf_counter() + for start, end in zip(starts, ends): + start.record() + operation() + end.record() + torch.cuda.synchronize() + wall_mean_us = (time.perf_counter() - wall_start) * 1e6 / repeats + elapsed_us = sorted(start.elapsed_time(end) * 1000.0 for start, end in zip(starts, ends)) + return { + "mean_us": statistics.fmean(elapsed_us), + "median_us": statistics.median(elapsed_us), + "p10_us": _percentile(elapsed_us, 0.10), + "p90_us": _percentile(elapsed_us, 0.90), + "p99_us": _percentile(elapsed_us, 0.99), + "min_us": elapsed_us[0], + "max_us": elapsed_us[-1], + "wall_mean_us": wall_mean_us, + "warmup": warmup, + "repeats": repeats, + } + + +def _torch_cumulative_lengths(torch, sequences: int, sequence_length: int, device) -> Any: + return torch.arange(sequences + 1, dtype=torch.int32, device=device).mul_(sequence_length) + + +def _cuda_impl_selected_block_indices(torch, context_length: int, concurrency: int, device) -> Any: + block_count = context_length // CUDA_IMPL_BLOCK_SIZE + forced_local_start = block_count - (CUDA_IMPL_LOCAL_BLOCKS + 1) + dynamic_count = CUDA_IMPL_SELECTED_BLOCKS - (CUDA_IMPL_LOCAL_BLOCKS + 1) - CUDA_IMPL_INIT_BLOCKS + selected = torch.cat( + ( + torch.arange(CUDA_IMPL_INIT_BLOCKS + dynamic_count, dtype=torch.int32, device=device), + torch.arange(forced_local_start, block_count, dtype=torch.int32, device=device), + ) + ) + if selected.numel() != CUDA_IMPL_SELECTED_BLOCKS or torch.unique(selected).numel() != CUDA_IMPL_SELECTED_BLOCKS: + raise RuntimeError("Equivalent workload did not produce the required unique block budget.") + return selected.reshape(1, 1, CUDA_IMPL_SELECTED_BLOCKS).expand(CUDA_IMPL_KV_HEADS, concurrency, -1).contiguous() + + +def _run_cuda_impl_scenario( + torch, + infllm_v2, + context_length: int, + concurrency: int, + warmup: int, + repeats: int, +) -> dict[str, Any]: + if context_length % CUDA_IMPL_BLOCK_SIZE != 0: + raise ValueError("context length must be divisible by block size") + if context_length // CUDA_IMPL_BLOCK_SIZE <= CUDA_IMPL_SELECTED_BLOCKS: + raise ValueError("context length must contain more blocks than the sparse selection budget") + device = torch.device("cuda:0") + dtype = torch.bfloat16 + torch.manual_seed(2026) + query = torch.randn( + concurrency, + CUDA_IMPL_QUERY_HEADS, + CUDA_IMPL_HEAD_DIM, + device=device, + dtype=dtype, + ) + key = torch.zeros( + concurrency * context_length, + CUDA_IMPL_KV_HEADS, + CUDA_IMPL_HEAD_DIM, + device=device, + dtype=dtype, + ) + value = torch.zeros_like(key) + fine_windows = (context_length - CUDA_IMPL_KERNEL_SIZE) // CUDA_IMPL_KERNEL_STRIDE + 1 + coarse_windows = (context_length - 4 * CUDA_IMPL_KERNEL_SIZE) // (4 * CUDA_IMPL_KERNEL_STRIDE) + 1 + compressed_fine = torch.zeros( + concurrency * fine_windows, + CUDA_IMPL_KV_HEADS, + CUDA_IMPL_HEAD_DIM, + device=device, + dtype=dtype, + ) + compressed_coarse = torch.zeros( + concurrency * coarse_windows, + CUDA_IMPL_KV_HEADS, + CUDA_IMPL_HEAD_DIM, + device=device, + dtype=dtype, + ) + cu_query = _torch_cumulative_lengths(torch, concurrency, 1, device) + cu_key = _torch_cumulative_lengths(torch, concurrency, context_length, device) + cu_fine = _torch_cumulative_lengths(torch, concurrency, fine_windows, device) + cu_coarse = _torch_cumulative_lengths(torch, concurrency, coarse_windows, device) + cache_lengths = torch.full((concurrency,), context_length - 1, dtype=torch.int32, device=device) + stage2_indices = _cuda_impl_selected_block_indices(torch, context_length, concurrency, device) + + def stage1_score(): + return infllm_v2.infllmv2_attn_stage1( + query, + compressed_fine, + compressed_coarse, + cu_seqlens_q=cu_query, + cu_seqlens_k=cu_fine, + cu_seqlens_v=cu_coarse, + max_seqlen_q=1, + max_seqlen_k=fine_windows, + causal=False, + ) + + def stage1_with_topk(): + scores = stage1_score() + pooled = infllm_v2.max_pooling_1d_varlen( + scores, + cu_query, + cu_fine, + cache_lengths, + 1, + fine_windows, + CUDA_IMPL_LOCAL_BLOCKS, + CUDA_IMPL_INIT_BLOCKS, + CUDA_IMPL_BLOCK_SIZE, + CUDA_IMPL_KERNEL_STRIDE, + ) + return pooled.topk(CUDA_IMPL_SELECTED_BLOCKS, dim=-1).indices.to(torch.int32).contiguous() + + def stage2_attention(indices=stage2_indices): + return infllm_v2.infllmv2_attn_varlen_func( + query, + key, + value, + cu_query, + cu_key, + 1, + context_length, + causal=False, + topk_idx=indices, + ) + + def sparse_chain(): + return stage2_attention(stage1_with_topk()) + + def dense_attention(): + return infllm_v2.infllmv2_attn_varlen_func( + query, + key, + value, + cu_query, + cu_key, + 1, + context_length, + causal=False, + ) + + scores = stage1_score() + if scores.shape != (CUDA_IMPL_KV_HEADS, concurrency, fine_windows): + raise RuntimeError(f"Unexpected Stage 1 score shape: {tuple(scores.shape)}") + selected = stage1_with_topk() + if selected.shape != ( + CUDA_IMPL_KV_HEADS, + concurrency, + CUDA_IMPL_SELECTED_BLOCKS, + ): + raise RuntimeError(f"Unexpected Top-K shape: {tuple(selected.shape)}") + stage2_output = stage2_attention() + if stage2_output.shape != query.shape: + raise RuntimeError(f"Unexpected Stage 2 output shape: {tuple(stage2_output.shape)}") + torch.cuda.synchronize() + + effective_tokens = concurrency * CUDA_IMPL_KV_HEADS * CUDA_IMPL_SELECTED_BLOCKS * CUDA_IMPL_BLOCK_SIZE + visible_tokens = concurrency * CUDA_IMPL_KV_HEADS * context_length + return { + "context_length": context_length, + "concurrency": concurrency, + "fine_windows": fine_windows, + "coarse_windows": coarse_windows, + "selected_blocks": CUDA_IMPL_SELECTED_BLOCKS, + "effective_stage2_kv_tokens": effective_tokens, + "visible_stage2_kv_tokens": visible_tokens, + "effective_sparse_kv_fraction": effective_tokens / visible_tokens, + "operators": { + "stage1_score": _time_torch_cuda(torch, stage1_score, warmup, repeats), + "stage1_with_topk": _time_torch_cuda(torch, stage1_with_topk, warmup, repeats), + "stage2_attention": _time_torch_cuda(torch, stage2_attention, warmup, repeats), + "sparse_decode_chain": _time_torch_cuda(torch, sparse_chain, warmup, repeats), + "dense_attention": _time_torch_cuda(torch, dense_attention, warmup, repeats), + }, + } + + +def _parse_cuda_impl_scenario(value: str) -> tuple[int, int]: + fields = value.split(":") + if len(fields) != 2: + raise argparse.ArgumentTypeError("scenario must be CONTEXT_LENGTH:CONCURRENCY") + try: + context_length, concurrency = (int(field) for field in fields) + except ValueError as exc: + raise argparse.ArgumentTypeError("scenario fields must be integers") from exc + if context_length < 1 or concurrency < 1: + raise argparse.ArgumentTypeError("scenario fields must be positive") + return context_length, concurrency + + +def run_cuda_impl_benchmark(args: argparse.Namespace) -> None: + _validate_output_args(args) + _require_int(args.gpu_index, "--gpu-index", 0) + _require_int(args.warmup, "--warmup", 1) + _require_int(args.repeats, "--repeats", 1) + repo = args.repo.resolve() + if not (repo / "infllm_v2").is_dir(): + raise FileNotFoundError(f"Missing infllm_v2 package under {repo}") + + sys.path.insert(0, str(repo)) + import infllm_v2 + import torch + from infllm_v2 import C + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + extension_path = Path(C.__file__).resolve() + extension_stat = extension_path.stat() + gpu_query = subprocess.run( + [ + "nvidia-smi", + f"--id={args.gpu_index}", + "--query-gpu=index,uuid,name,memory.total", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + git_head = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repo, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + scenarios = args.cuda_impl_scenarios or [(32768, 1), (131072, 4)] + if len(set(scenarios)) != len(scenarios): + raise ValueError("--scenario entries must be unique") + results = [ + _run_cuda_impl_scenario(torch, infllm_v2, context, concurrency, args.warmup, args.repeats) + for context, concurrency in scenarios + ] + output = { + "kind": "infllmv2_cuda_impl_operator_benchmark", + "repo": str(repo), + "git_head": git_head, + "gpu": gpu_query, + "torch_version": torch.__version__, + "dtype": "bfloat16", + "model_shape": { + "query_heads": CUDA_IMPL_QUERY_HEADS, + "kv_heads": CUDA_IMPL_KV_HEADS, + "gqa_group_size": CUDA_IMPL_QUERY_HEADS // CUDA_IMPL_KV_HEADS, + "head_dim": CUDA_IMPL_HEAD_DIM, + "block_size": CUDA_IMPL_BLOCK_SIZE, + "selected_blocks": CUDA_IMPL_SELECTED_BLOCKS, + }, + "loaded_extension": { + "path": str(extension_path), + "mtime_ns": extension_stat.st_mtime_ns, + "mtime_utc": datetime.datetime.fromtimestamp(extension_stat.st_mtime, datetime.timezone.utc).isoformat(), + "size_bytes": extension_stat.st_size, + "sha256": _sha256_file(extension_path), + }, + "scenarios": results, + } + _write_json(args.output, output, args.overwrite) + for scenario in results: + operators = scenario["operators"] + print( + f"{scenario['context_length']} tokens / concurrency {scenario['concurrency']}: " + f"stage1+topk={operators['stage1_with_topk']['median_us']:.3f} us, " + f"stage2={operators['stage2_attention']['median_us']:.3f} us, " + f"chain={operators['sparse_decode_chain']['median_us']:.3f} us, " + f"dense={operators['dense_attention']['median_us']:.3f} us" + ) + print(f"Saved local cuda_impl benchmark to {args.output}") + + +def _add_output_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--output", + required=True, + type=Path, + help="JSON output path; its parent directory must exist.", + ) + parser.add_argument("--overwrite", action="store_true", help="Replace --output if it already exists.") + + +def _add_workload_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--base-url", + required=True, + help="OpenAI-compatible server base URL, e.g. http://127.0.0.1:8000.", + ) + parser.add_argument("--endpoint", default="/v1/chat/completions") + parser.add_argument("--model", required=True, help="Model name accepted by the server.") + parser.add_argument( + "--tokenizer", + help="Optional tokenizer name/path forwarded to benchmark_serving.py.", + ) + parser.add_argument("--seed", default=2026, type=int) + parser.add_argument("--num-prompts", default=1, type=int) + parser.add_argument("--input-len", required=True, type=int) + parser.add_argument("--output-len", default=128, type=int) + parser.add_argument("--request-rate", default=float("inf"), type=float) + parser.add_argument("--max-concurrency", default=1, type=int) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + subparsers = parser.add_subparsers(dest="command", required=True) + + run_parser = subparsers.add_parser( + "run", + help="Run benchmark_serving.py against an existing OpenAI-compatible server and sample GPU memory.", + ) + run_parser.add_argument("--variant", required=True, choices=("dense", "sparse")) + _add_workload_arguments(run_parser) + run_parser.add_argument("--gpu-index", required=True, type=int, help="GPU index sampled with nvidia-smi.") + run_parser.add_argument( + "--sample-interval", + default=0.1, + type=float, + help="nvidia-smi polling interval in seconds.", + ) + run_parser.add_argument( + "--selector-diagnostic", + type=Path, + help=("Bound sparse-diagnostic JSON; required for sparse timing runs and rejected for dense runs."), + ) + _add_output_arguments(run_parser) + run_parser.set_defaults(handler=run_benchmark) + + diagnostic_parser = subparsers.add_parser( + "sparse-diagnostic", + help=( + "Run an un-warmed sparse request workload and bind a newly generated selector trace to its exact prompts." + ), + ) + _add_workload_arguments(diagnostic_parser) + diagnostic_parser.add_argument( + "--trace-path", + required=True, + type=Path, + help=(f"New selector trace path configured on the sparse server through {SELECTOR_TRACE_PATH_ENV}."), + ) + _add_output_arguments(diagnostic_parser) + diagnostic_parser.set_defaults(handler=run_selector_diagnostic) + + report_parser = subparsers.add_parser("report", help="Compare matching dense and sparse run JSON files.") + report_parser.add_argument("--dense-result", required=True, type=Path) + report_parser.add_argument("--sparse-result", required=True, type=Path) + _add_output_arguments(report_parser) + report_parser.set_defaults(handler=create_report) + + selector_parser = subparsers.add_parser( + "selector", + help="Inspect a backend selector trace without binding it to a serving workload.", + ) + selector_parser.add_argument( + "--samples", + required=True, + type=Path, + help="Backend infllmv2_selector_samples trace JSON.", + ) + _add_output_arguments(selector_parser) + selector_parser.set_defaults(handler=summarize_selector) + + operator_parser = subparsers.add_parser( + "operators", + help=("Time InfLLM-V2 update, Stage 1, Stage 2, and their decode chain with CUDA events."), + ) + operator_parser.add_argument( + "--scenario", + dest="scenarios", + action="append", + type=_parse_operator_scenario, + help=("CONTEXT_LENGTH:CONCURRENCY; repeat for a matrix. Defaults to 32768:1 and 131072:4."), + ) + operator_parser.add_argument("--device", default="gpu:0", help="Paddle CUDA device visible to this process.") + operator_parser.add_argument( + "--gpu-index", + required=True, + type=int, + help="Physical GPU index reported through nvidia-smi.", + ) + operator_parser.add_argument( + "--dtype", + default="bfloat16", + choices=("float16", "bfloat16", "float32"), + ) + operator_parser.add_argument("--seed", default=2026, type=int) + operator_parser.add_argument("--warmup", default=20, type=int) + operator_parser.add_argument("--repeats", default=100, type=int) + _add_output_arguments(operator_parser) + operator_parser.set_defaults(handler=run_operator_benchmark) + + prefill_parser = subparsers.add_parser( + "prefill", + help="Compare complete dense and InfLLM-V2 sparse prefill paths with CUDA events.", + ) + prefill_parser.add_argument( + "--context-length", + action="append", + type=int, + dest="context_lengths", + help="Repeat for multiple lengths. Defaults to 16384 and 32768.", + ) + prefill_parser.add_argument("--device", default="gpu:0") + prefill_parser.add_argument("--dtype", choices=("float16", "bfloat16"), default="bfloat16") + prefill_parser.add_argument("--seed", type=int, default=2026) + prefill_parser.add_argument("--warmup", type=int, default=10) + prefill_parser.add_argument("--repeats", type=int, default=30) + _add_output_arguments(prefill_parser) + prefill_parser.set_defaults(handler=run_prefill_benchmark) + + cuda_impl_parser = subparsers.add_parser( + "cuda-impl", + help="Benchmark a local infllm_v2 PyTorch extension with an equivalent decode workload.", + ) + cuda_impl_parser.add_argument("--repo", required=True, type=Path) + cuda_impl_parser.add_argument("--gpu-index", required=True, type=int) + cuda_impl_parser.add_argument( + "--scenario", + dest="cuda_impl_scenarios", + action="append", + type=_parse_cuda_impl_scenario, + help="CONTEXT_LENGTH:CONCURRENCY; defaults to 32768:1 and 131072:4.", + ) + cuda_impl_parser.add_argument("--warmup", default=20, type=int) + cuda_impl_parser.add_argument("--repeats", default=100, type=int) + _add_output_arguments(cuda_impl_parser) + cuda_impl_parser.set_defaults(handler=run_cuda_impl_benchmark) + return parser + + +def main() -> None: + args = build_parser().parse_args() + args.handler(args) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/yaml/minicpm41-8b-32k-bf16-tp1.yaml b/benchmarks/yaml/minicpm41-8b-32k-bf16-tp1.yaml new file mode 100644 index 00000000000..a33b2d1b826 --- /dev/null +++ b/benchmarks/yaml/minicpm41-8b-32k-bf16-tp1.yaml @@ -0,0 +1,4 @@ +max_model_len: 32768 +max_num_seqs: 8 +kv_cache_ratio: 0.75 +tensor_parallel_size: 1 diff --git a/benchmarks/yaml/minicpm41-8b-32k-wint4-tp1.yaml b/benchmarks/yaml/minicpm41-8b-32k-wint4-tp1.yaml new file mode 100644 index 00000000000..6858efb9902 --- /dev/null +++ b/benchmarks/yaml/minicpm41-8b-32k-wint4-tp1.yaml @@ -0,0 +1,5 @@ +max_model_len: 32768 +max_num_seqs: 8 +kv_cache_ratio: 0.75 +tensor_parallel_size: 1 +quantization: wint4 diff --git a/benchmarks/yaml/minicpm41-8b-32k-wint8-tp1.yaml b/benchmarks/yaml/minicpm41-8b-32k-wint8-tp1.yaml new file mode 100644 index 00000000000..8bdd97da493 --- /dev/null +++ b/benchmarks/yaml/minicpm41-8b-32k-wint8-tp1.yaml @@ -0,0 +1,5 @@ +max_model_len: 32768 +max_num_seqs: 8 +kv_cache_ratio: 0.75 +tensor_parallel_size: 1 +quantization: wint8 diff --git a/custom_ops/gpu_ops/cpp_extensions.cc b/custom_ops/gpu_ops/cpp_extensions.cc index 911695e7412..0f294aa0b78 100644 --- a/custom_ops/gpu_ops/cpp_extensions.cc +++ b/custom_ops/gpu_ops/cpp_extensions.cc @@ -64,6 +64,57 @@ paddle::Tensor CustomNumpyToTensor(py::array numpy_array, return tensor; } +std::vector InfLLMV2UpdateCompressedK( + const paddle::Tensor& current_tokens, + const paddle::Tensor& key_cache, + paddle::Tensor& compressed_k, + paddle::Tensor& compressed_k2, + const paddle::Tensor& block_tables, + const paddle::Tensor& seq_lens_decoder, + const paddle::Tensor& seq_lens_this_time, + const paddle::Tensor& batch_id_per_token, + const paddle::Tensor& cu_seqlens_q, + int kernel_size, + int kernel_stride); + +std::vector InfLLMV2SelectBlocks( + const paddle::Tensor& query, + const paddle::Tensor& compressed_k, + const paddle::Tensor& compressed_k2, + const paddle::Tensor& block_tables, + const paddle::Tensor& seq_lens_decoder, + const paddle::Tensor& seq_lens_this_time, + const paddle::Tensor& batch_id_per_token, + const paddle::Tensor& cu_seqlens_q, + paddle::Tensor& topk_indices, + paddle::Tensor& block_scores, + paddle::Tensor& selected_counts, + paddle::Tensor& coarse_lse, + paddle::Tensor& coarse_partial_max, + paddle::Tensor& coarse_partial_sum, + int block_size, + int kernel_size, + int kernel_stride, + int topk, + int dense_len, + int init_blocks, + int local_blocks); + +std::vector InfLLMV2AttentionForward( + const paddle::Tensor& query, + const paddle::Tensor& key_cache, + const paddle::Tensor& value_cache, + const paddle::Tensor& block_tables, + const paddle::Tensor& seq_lens_decoder, + const paddle::Tensor& seq_lens_this_time, + const paddle::Tensor& batch_id_per_token, + const paddle::Tensor& cu_seqlens_q, + const paddle::Tensor& topk_indices, + paddle::Tensor& out, + paddle::Tensor& partial_acc, + paddle::Tensor& partial_max, + paddle::Tensor& partial_sum); + void FlashAttentionMask(const paddle::Tensor& q_input, const paddle::Tensor& k_input, const paddle::Tensor& v_input, @@ -2009,6 +2060,15 @@ PYBIND11_MODULE(fastdeploy_ops, m) { "reasoning_phase_token_constraint function"); m.def("get_attn_mask_q", &get_attn_mask_q, "get_attn_mask_q function"); + m.def("infllmv2_update_compressed_k", + &InfLLMV2UpdateCompressedK, + "infllmv2_update_compressed_k function"); + m.def("infllmv2_select_blocks", + &InfLLMV2SelectBlocks, + "infllmv2_select_blocks function"); + m.def("infllmv2_attention_forward", + &InfLLMV2AttentionForward, + "infllmv2_attention_forward function"); m.def("custom_numpy_to_tensor", &CustomNumpyToTensor, diff --git a/custom_ops/gpu_ops/infllmv2_attention/infllmv2.cu b/custom_ops/gpu_ops/infllmv2_attention/infllmv2.cu new file mode 100644 index 00000000000..0cc23fb7ee2 --- /dev/null +++ b/custom_ops/gpu_ops/infllmv2_attention/infllmv2.cu @@ -0,0 +1,1143 @@ +// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "paddle/extension.h" +#include "infllmv2_attention/infllmv2_impl.cuh" + +#include +#include +#include +#include +#include +#include + +#ifndef PD_BUILD_STATIC_OP +#define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name) +#endif + +namespace { + +template +struct InfLLMTypeTraits; + +template <> +struct InfLLMTypeTraits { + using NativeType = float; + using PaddleType = float; +}; + +template <> +struct InfLLMTypeTraits { + using NativeType = half; + using PaddleType = phi::dtype::float16; +}; + +template <> +struct InfLLMTypeTraits { + using NativeType = __nv_bfloat16; + using PaddleType = phi::dtype::bfloat16; +}; + +void CheckCUDA(const paddle::Tensor& tensor, const char* name) { + PD_CHECK(tensor.is_gpu(), "InfLLM-V2 ", name, " must be a CUDA tensor."); +} + +void CheckInt32(const paddle::Tensor& tensor, const char* name) { + PD_CHECK(tensor.dtype() == paddle::DataType::INT32, + "InfLLM-V2 ", + name, + " must use int32 metadata."); +} + +void CheckFloatingType(const paddle::Tensor& tensor, const char* name) { + const auto dtype = tensor.dtype(); + PD_CHECK(dtype == paddle::DataType::FLOAT32 || + dtype == paddle::DataType::FLOAT16 || + dtype == paddle::DataType::BFLOAT16, + "InfLLM-V2 ", + name, + " only supports fp32, fp16, or bf16."); +} + +void CheckCommonMetadata(const paddle::Tensor& block_tables, + const paddle::Tensor& seq_lens_decoder, + const paddle::Tensor& seq_lens_this_time, + const paddle::Tensor& batch_id_per_token, + const paddle::Tensor& cu_seqlens_q, + int tokens) { + for (const auto* tensor : {&block_tables, + &seq_lens_decoder, + &seq_lens_this_time, + &batch_id_per_token, + &cu_seqlens_q}) { + CheckCUDA(*tensor, "runtime metadata"); + CheckInt32(*tensor, "runtime metadata"); + } + PD_CHECK(block_tables.shape().size() == 2, + "InfLLM-V2 block_tables must be rank 2."); + PD_CHECK(seq_lens_decoder.shape().size() == 1 && + seq_lens_this_time.shape().size() == 1 && + batch_id_per_token.shape().size() == 1 && + cu_seqlens_q.shape().size() == 1, + "InfLLM-V2 sequence metadata must be rank 1."); + const int64_t batch_size = block_tables.shape()[0]; + PD_CHECK(batch_size > 0 && block_tables.shape()[1] > 0, + "InfLLM-V2 block_tables dimensions must be positive."); + PD_CHECK( + seq_lens_decoder.shape()[0] == batch_size && + seq_lens_this_time.shape()[0] == batch_size, + "InfLLM-V2 sequence-length tensors must match block_tables batch size."); + PD_CHECK(cu_seqlens_q.shape()[0] == batch_size + 1, + "InfLLM-V2 cu_seqlens_q must have batch_size + 1 elements."); + PD_CHECK(batch_id_per_token.shape()[0] == tokens, + "InfLLM-V2 batch_id_per_token length must match query tokens."); +} + +template +void LaunchUpdateCompressedK(const paddle::Tensor& current_tokens, + const paddle::Tensor& key_cache, + paddle::Tensor& compressed_k, + paddle::Tensor& compressed_k2, + const paddle::Tensor& block_tables, + const paddle::Tensor& seq_lens_decoder, + const paddle::Tensor& batch_id_per_token, + const paddle::Tensor& cu_seqlens_q, + int kernel_size, + int kernel_stride) { + using Traits = InfLLMTypeTraits; + using NativeT = typename Traits::NativeType; + using PaddleT = typename Traits::PaddleType; + const int tokens = static_cast(current_tokens.shape()[0]); + const int kv_heads = static_cast(key_cache.shape()[1]); + const int head_dim = static_cast(key_cache.shape()[3]); + if (tokens == 0) { + return; + } + dim3 grid(tokens, kv_heads, 2); + const int threads = std::min(256, std::max(32, head_dim)); + fastdeploy::InfLLMV2UpdateCompressedKKernel + <<>>( + reinterpret_cast(key_cache.data()), + reinterpret_cast(compressed_k.data()), + reinterpret_cast(compressed_k2.data()), + block_tables.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + tokens, + static_cast(block_tables.shape()[0]), + static_cast(key_cache.shape()[0]), + static_cast(block_tables.shape()[1]), + kv_heads, + static_cast(key_cache.shape()[2]), + head_dim, + kernel_size, + kernel_stride); +} + +template +void LaunchSelectBlocks(const paddle::Tensor& query, + const paddle::Tensor& compressed_k, + const paddle::Tensor& compressed_k2, + const paddle::Tensor& block_tables, + const paddle::Tensor& seq_lens_decoder, + const paddle::Tensor& batch_id_per_token, + const paddle::Tensor& cu_seqlens_q, + paddle::Tensor& topk_indices, + paddle::Tensor& block_scores, + paddle::Tensor& selected_counts, + paddle::Tensor& coarse_lse, + paddle::Tensor& coarse_partial_max, + paddle::Tensor& coarse_partial_sum, + int block_size, + int kernel_size, + int kernel_stride, + int topk, + int dense_len, + int init_blocks, + int local_blocks) { + using Traits = InfLLMTypeTraits; + using NativeT = typename Traits::NativeType; + using PaddleT = typename Traits::PaddleType; + const int tokens = static_cast(query.shape()[0]); + const int query_heads = static_cast(query.shape()[1]); + const int head_dim = static_cast(query.shape()[2]); + const int kv_heads = static_cast(compressed_k.shape()[1]); + const int max_blocks = static_cast(block_tables.shape()[1]); + const int coarse_splits = static_cast(coarse_partial_max.shape()[2]); + if (tokens == 0) { + return; + } + const int group_size = query_heads / kv_heads; + constexpr int coarse_threads = 256; + if constexpr (!std::is_same_v) { + if (group_size == 16 && head_dim % 16 == 0) { + constexpr int tensor_tile = 16; + const size_t tensor_shared_k_bytes = + static_cast(tensor_tile) * head_dim * sizeof(NativeT); + const size_t tensor_scores_offset = + (tensor_shared_k_bytes + alignof(float) - 1) & ~(alignof(float) - 1); + const size_t tensor_shared_bytes = + tensor_scores_offset + tensor_tile * tensor_tile * sizeof(float); + const dim3 tensor_coarse_grid(tokens, kv_heads, coarse_splits); + fastdeploy::InfLLMV2CoarseLSETensorCoreSplitKernel + <<>>( + reinterpret_cast(query.data()), + reinterpret_cast(compressed_k2.data()), + block_tables.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + coarse_partial_max.data(), + coarse_partial_sum.data(), + tokens, + static_cast(block_tables.shape()[0]), + max_blocks, + query_heads, + kv_heads, + block_size, + head_dim, + kernel_size, + kernel_stride, + coarse_splits); + constexpr int combine_threads = 128; + const dim3 combine_grid(tokens, query_heads); + fastdeploy:: + InfLLMV2CoarseLSECombineKernel<<>>( + coarse_partial_max.data(), + coarse_partial_sum.data(), + batch_id_per_token.data(), + coarse_lse.data(), + tokens, + static_cast(block_tables.shape()[0]), + query_heads, + coarse_splits); + } else { + const dim3 coarse_grid(tokens, query_heads); + fastdeploy::InfLLMV2CoarseLSEKernel + <<>>( + reinterpret_cast(query.data()), + reinterpret_cast(compressed_k2.data()), + block_tables.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + coarse_lse.data(), + coarse_partial_max.data(), + coarse_partial_sum.data(), + tokens, + static_cast(block_tables.shape()[0]), + max_blocks, + query_heads, + kv_heads, + block_size, + head_dim, + kernel_size, + kernel_stride, + coarse_splits); + } + } else { + const dim3 coarse_grid(tokens, query_heads); + fastdeploy::InfLLMV2CoarseLSEKernel + <<>>( + reinterpret_cast(query.data()), + reinterpret_cast(compressed_k2.data()), + block_tables.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + coarse_lse.data(), + coarse_partial_max.data(), + coarse_partial_sum.data(), + tokens, + static_cast(block_tables.shape()[0]), + max_blocks, + query_heads, + kv_heads, + block_size, + head_dim, + kernel_size, + kernel_stride, + coarse_splits); + } + + constexpr int score_blocks_per_cta = 4; + const dim3 score_grid( + tokens, + kv_heads, + (max_blocks + score_blocks_per_cta - 1) / score_blocks_per_cta); + const int score_windows = block_size / kernel_stride + 1; + const int score_threads = group_size * 32; + const size_t shared_k_bytes = + static_cast(score_windows) * head_dim * sizeof(NativeT); + const size_t score_shared_bytes = + ((shared_k_bytes + alignof(float) - 1) & ~(alignof(float) - 1)) + + static_cast(score_windows) * group_size * sizeof(float); + if constexpr (!std::is_same_v) { + if (group_size == 16 && head_dim % 16 == 0 && score_windows <= 16) { + constexpr int tensor_score_warps = 4; + constexpr int tensor_tile = 16; + const size_t tensor_shared_k_bytes = + static_cast(tensor_tile) * head_dim * sizeof(NativeT); + const size_t tensor_scores_offset = + (tensor_shared_k_bytes + alignof(float) - 1) & ~(alignof(float) - 1); + const size_t tensor_warp_bytes = + tensor_scores_offset + tensor_tile * tensor_tile * sizeof(float); + fastdeploy::InfLLMV2BlockScoreTensorCoreKernel + <<>>( + reinterpret_cast(query.data()), + reinterpret_cast(compressed_k.data()), + block_tables.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + coarse_lse.data(), + block_scores.data(), + tokens, + static_cast(block_tables.shape()[0]), + max_blocks, + query_heads, + kv_heads, + block_size, + head_dim, + kernel_size, + kernel_stride, + init_blocks, + local_blocks); + } else { + fastdeploy::InfLLMV2BlockScoreKernel + <<>>( + reinterpret_cast(query.data()), + reinterpret_cast(compressed_k.data()), + block_tables.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + coarse_lse.data(), + block_scores.data(), + tokens, + static_cast(block_tables.shape()[0]), + max_blocks, + query_heads, + kv_heads, + block_size, + head_dim, + kernel_size, + kernel_stride, + init_blocks, + local_blocks); + } + } else { + fastdeploy::InfLLMV2BlockScoreKernel + <<>>( + reinterpret_cast(query.data()), + reinterpret_cast(compressed_k.data()), + block_tables.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + coarse_lse.data(), + block_scores.data(), + tokens, + static_cast(block_tables.shape()[0]), + max_blocks, + query_heads, + kv_heads, + block_size, + head_dim, + kernel_size, + kernel_stride, + init_blocks, + local_blocks); + } + + const dim3 topk_grid(tokens, kv_heads); + fastdeploy::InfLLMV2TopKKernel<<>>( + block_scores.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + topk_indices.data(), + selected_counts.data(), + tokens, + static_cast(block_tables.shape()[0]), + kv_heads, + max_blocks, + static_cast(topk_indices.shape()[2]), + block_size, + topk, + dense_len, + local_blocks); +} + +template +void LaunchSparseAttention(const paddle::Tensor& query, + const paddle::Tensor& key_cache, + const paddle::Tensor& value_cache, + const paddle::Tensor& block_tables, + const paddle::Tensor& seq_lens_decoder, + const paddle::Tensor& batch_id_per_token, + const paddle::Tensor& cu_seqlens_q, + const paddle::Tensor& topk_indices, + paddle::Tensor& out, + paddle::Tensor& partial_acc, + paddle::Tensor& partial_max, + paddle::Tensor& partial_sum) { + using Traits = InfLLMTypeTraits; + using NativeT = typename Traits::NativeType; + using PaddleT = typename Traits::PaddleType; + const int tokens = static_cast(query.shape()[0]); + const int query_heads = static_cast(query.shape()[1]); + const int head_dim = static_cast(query.shape()[2]); + const int capacity = static_cast(topk_indices.shape()[2]); + const int block_size = static_cast(key_cache.shape()[2]); + if (tokens == 0) { + return; + } + const int kv_heads = static_cast(key_cache.shape()[1]); + const int group_size = query_heads / kv_heads; + const int splits = static_cast(partial_acc.shape()[2]); + const int blocks_per_split = (capacity + splits - 1) / splits; + const int split_threads = group_size * 32; + const size_t split_shared_bytes = + static_cast(2 * 8 * head_dim) * sizeof(NativeT); + const dim3 split_grid(tokens, kv_heads, splits); + if constexpr (!std::is_same_v) { + if (group_size == 16 && head_dim % 16 == 0) { + constexpr int tensor_tile = 16; + constexpr int tensor_warps = 4; + const size_t tensor_values_bytes = + static_cast(2 * tensor_tile * head_dim) * sizeof(NativeT); + const size_t tensor_scores_offset = + (tensor_values_bytes + alignof(float) - 1) & ~(alignof(float) - 1); + const size_t tensor_shared_bytes = + tensor_scores_offset + group_size * tensor_tile * sizeof(float) + + group_size * tensor_tile * sizeof(NativeT) + + group_size * head_dim * sizeof(float) + + tensor_warps * group_size * tensor_tile * sizeof(float) + + 3 * group_size * sizeof(float); + fastdeploy::InfLLMV2SparseAttentionTensorCoreSplitKVKernel + <<>>( + reinterpret_cast(query.data()), + reinterpret_cast(key_cache.data()), + reinterpret_cast(value_cache.data()), + block_tables.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + topk_indices.data(), + partial_acc.data(), + partial_max.data(), + partial_sum.data(), + tokens, + static_cast(block_tables.shape()[0]), + static_cast(key_cache.shape()[0]), + static_cast(block_tables.shape()[1]), + query_heads, + kv_heads, + block_size, + head_dim, + capacity, + splits, + blocks_per_split); + } else { + fastdeploy::InfLLMV2SparseAttentionSplitKVKernel + <<>>( + reinterpret_cast(query.data()), + reinterpret_cast(key_cache.data()), + reinterpret_cast(value_cache.data()), + block_tables.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + topk_indices.data(), + partial_acc.data(), + partial_max.data(), + partial_sum.data(), + tokens, + static_cast(block_tables.shape()[0]), + static_cast(key_cache.shape()[0]), + static_cast(block_tables.shape()[1]), + query_heads, + kv_heads, + block_size, + head_dim, + capacity, + splits, + blocks_per_split); + } + } else { + fastdeploy::InfLLMV2SparseAttentionSplitKVKernel + <<>>( + reinterpret_cast(query.data()), + reinterpret_cast(key_cache.data()), + reinterpret_cast(value_cache.data()), + block_tables.data(), + seq_lens_decoder.data(), + batch_id_per_token.data(), + cu_seqlens_q.data(), + topk_indices.data(), + partial_acc.data(), + partial_max.data(), + partial_sum.data(), + tokens, + static_cast(block_tables.shape()[0]), + static_cast(key_cache.shape()[0]), + static_cast(block_tables.shape()[1]), + query_heads, + kv_heads, + block_size, + head_dim, + capacity, + splits, + blocks_per_split); + } + constexpr int combine_threads = 256; + const dim3 combine_grid(tokens, query_heads); + fastdeploy::InfLLMV2SparseAttentionCombineKernel + <<>>(partial_acc.data(), + partial_max.data(), + partial_sum.data(), + reinterpret_cast(out.data()), + tokens, + query_heads, + head_dim, + splits); +} + +} // namespace + +std::vector InfLLMV2UpdateCompressedK( + const paddle::Tensor& current_tokens, + const paddle::Tensor& key_cache, + paddle::Tensor& compressed_k, + paddle::Tensor& compressed_k2, + const paddle::Tensor& block_tables, + const paddle::Tensor& seq_lens_decoder, + const paddle::Tensor& seq_lens_this_time, + const paddle::Tensor& batch_id_per_token, + const paddle::Tensor& cu_seqlens_q, + int kernel_size, + int kernel_stride) { + CheckCUDA(current_tokens, "current_tokens"); + CheckCUDA(key_cache, "key_cache"); + CheckCUDA(compressed_k, "compressed_k"); + CheckCUDA(compressed_k2, "compressed_k2"); + CheckFloatingType(key_cache, "key_cache"); + PD_CHECK(current_tokens.dtype() == key_cache.dtype() && + compressed_k.dtype() == key_cache.dtype() && + compressed_k2.dtype() == key_cache.dtype(), + "InfLLM-V2 update tensors must have matching dtypes."); + PD_CHECK( + current_tokens.shape().size() == 2 || current_tokens.shape().size() == 3, + "InfLLM-V2 current_tokens must be rank 2 or rank 3."); + PD_CHECK(key_cache.shape().size() == 4, + "InfLLM-V2 key_cache must have shape " + "[physical_blocks, kv_heads, block_size, head_dim]."); + PD_CHECK( + compressed_k.shape().size() == 4 && compressed_k2.shape().size() == 4, + "InfLLM-V2 compressed K caches must be rank 4."); + PD_CHECK(kernel_size > 0 && kernel_stride > 0, + "InfLLM-V2 kernel_size and kernel_stride must be positive."); + const int64_t block_size = key_cache.shape()[2]; + PD_CHECK( + block_size % kernel_stride == 0 && block_size % (4 * kernel_stride) == 0, + "InfLLM-V2 block_size must be divisible by kernel_stride and " + "4 * kernel_stride."); + const std::vector expected_fine = {key_cache.shape()[0], + key_cache.shape()[1], + block_size / kernel_stride, + key_cache.shape()[3]}; + const std::vector expected_coarse = { + key_cache.shape()[0], + key_cache.shape()[1], + block_size / (4 * kernel_stride), + key_cache.shape()[3]}; + PD_CHECK( + compressed_k.shape() == expected_fine, + "InfLLM-V2 compressed_k shape does not match the fine summary layout."); + PD_CHECK(compressed_k2.shape() == expected_coarse, + "InfLLM-V2 compressed_k2 shape does not match the coarse summary " + "layout."); + CheckCommonMetadata(block_tables, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + static_cast(current_tokens.shape()[0])); + + switch (key_cache.dtype()) { + case paddle::DataType::FLOAT32: + LaunchUpdateCompressedK(current_tokens, + key_cache, + compressed_k, + compressed_k2, + block_tables, + seq_lens_decoder, + batch_id_per_token, + cu_seqlens_q, + kernel_size, + kernel_stride); + break; + case paddle::DataType::FLOAT16: + LaunchUpdateCompressedK(current_tokens, + key_cache, + compressed_k, + compressed_k2, + block_tables, + seq_lens_decoder, + batch_id_per_token, + cu_seqlens_q, + kernel_size, + kernel_stride); + break; + case paddle::DataType::BFLOAT16: + LaunchUpdateCompressedK(current_tokens, + key_cache, + compressed_k, + compressed_k2, + block_tables, + seq_lens_decoder, + batch_id_per_token, + cu_seqlens_q, + kernel_size, + kernel_stride); + break; + default: + PD_THROW("InfLLM-V2 update only supports fp32, fp16, or bf16."); + } + return {compressed_k, compressed_k2}; +} + +std::vector InfLLMV2SelectBlocks( + const paddle::Tensor& query, + const paddle::Tensor& compressed_k, + const paddle::Tensor& compressed_k2, + const paddle::Tensor& block_tables, + const paddle::Tensor& seq_lens_decoder, + const paddle::Tensor& seq_lens_this_time, + const paddle::Tensor& batch_id_per_token, + const paddle::Tensor& cu_seqlens_q, + paddle::Tensor& topk_indices, + paddle::Tensor& block_scores, + paddle::Tensor& selected_counts, + paddle::Tensor& coarse_lse, + paddle::Tensor& coarse_partial_max, + paddle::Tensor& coarse_partial_sum, + int block_size, + int kernel_size, + int kernel_stride, + int topk, + int dense_len, + int init_blocks, + int local_blocks) { + for (const paddle::Tensor* tensor : + std::vector{&query, + &compressed_k, + &compressed_k2, + &topk_indices, + &block_scores, + &selected_counts, + &coarse_lse, + &coarse_partial_max, + &coarse_partial_sum}) { + CheckCUDA(*tensor, "Stage 1 tensor"); + } + CheckFloatingType(query, "query"); + PD_CHECK(query.shape().size() == 3, + "InfLLM-V2 query must have shape [tokens, heads, head_dim]."); + PD_CHECK( + compressed_k.shape().size() == 4 && compressed_k2.shape().size() == 4, + "InfLLM-V2 compressed K caches must be rank 4."); + PD_CHECK(query.dtype() == compressed_k.dtype() && + query.dtype() == compressed_k2.dtype(), + "InfLLM-V2 query and compressed K dtypes must match."); + PD_CHECK(block_size > 0 && kernel_size > 0 && kernel_stride > 0 && topk > 0 && + dense_len >= 0 && init_blocks >= 0 && local_blocks >= 0, + "InfLLM-V2 Stage 1 attributes are outside their valid ranges."); + PD_CHECK( + block_size % kernel_stride == 0 && block_size % (4 * kernel_stride) == 0, + "InfLLM-V2 block_size must be divisible by kernel_stride and " + "4 * kernel_stride."); + const int64_t tokens = query.shape()[0]; + const int64_t query_heads = query.shape()[1]; + const int64_t head_dim = query.shape()[2]; + const int64_t kv_heads = compressed_k.shape()[1]; + const int64_t max_blocks = block_tables.shape()[1]; + PD_CHECK(kv_heads > 0 && query_heads % kv_heads == 0, + "InfLLM-V2 query heads must be divisible by KV heads."); + PD_CHECK(query_heads / kv_heads <= 32, + "InfLLM-V2 Stage 1 supports at most 32 query heads per KV head."); + PD_CHECK(head_dim > 0 && head_dim <= 256 && + head_dim == compressed_k.shape()[3] && + head_dim == compressed_k2.shape()[3], + "InfLLM-V2 query and compressed K head dimensions must match."); + PD_CHECK(compressed_k.shape()[2] == block_size / kernel_stride && + compressed_k2.shape()[2] == block_size / (4 * kernel_stride), + "InfLLM-V2 compressed K semantic strides are incompatible with " + "block_size."); + PD_CHECK(topk_indices.dtype() == paddle::DataType::INT32 && + selected_counts.dtype() == paddle::DataType::INT32, + "InfLLM-V2 topk_indices and selected_counts must be int32."); + PD_CHECK(block_scores.dtype() == paddle::DataType::FLOAT32 && + coarse_lse.dtype() == paddle::DataType::FLOAT32 && + coarse_partial_max.dtype() == paddle::DataType::FLOAT32 && + coarse_partial_sum.dtype() == paddle::DataType::FLOAT32, + "InfLLM-V2 Stage 1 score workspaces must be float32."); + PD_CHECK(topk_indices.shape().size() == 3 && + topk_indices.shape()[0] == tokens && + topk_indices.shape()[1] == kv_heads, + "InfLLM-V2 topk_indices must have shape " + "[tokens, kv_heads, capacity]."); + PD_CHECK(block_scores.shape() == + std::vector({tokens, kv_heads, max_blocks}), + "InfLLM-V2 block_scores shape is invalid."); + PD_CHECK(selected_counts.shape() == std::vector({tokens, kv_heads}), + "InfLLM-V2 selected_counts shape is invalid."); + PD_CHECK(coarse_lse.shape() == std::vector({tokens, query_heads}), + "InfLLM-V2 coarse_lse shape is invalid."); + PD_CHECK(coarse_partial_max.shape().size() == 3 && + coarse_partial_max.shape() == coarse_partial_sum.shape() && + coarse_partial_max.shape()[0] == tokens && + coarse_partial_max.shape()[1] == query_heads && + coarse_partial_max.shape()[2] > 0, + "InfLLM-V2 coarse partial workspace shape is invalid."); + const int64_t required_capacity = std::min( + max_blocks, + std::max(topk + local_blocks, + (dense_len + block_size - 1) / block_size)); + PD_CHECK(topk_indices.shape()[2] >= required_capacity, + "InfLLM-V2 topk_indices capacity is too small for sparse and " + "select-all modes."); + PD_CHECK(max_blocks <= 2048, + "InfLLM-V2 Stage 1 supports at most 2048 candidate blocks."); + CheckCommonMetadata(block_tables, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + static_cast(tokens)); + + switch (query.dtype()) { + case paddle::DataType::FLOAT32: + LaunchSelectBlocks(query, + compressed_k, + compressed_k2, + block_tables, + seq_lens_decoder, + batch_id_per_token, + cu_seqlens_q, + topk_indices, + block_scores, + selected_counts, + coarse_lse, + coarse_partial_max, + coarse_partial_sum, + block_size, + kernel_size, + kernel_stride, + topk, + dense_len, + init_blocks, + local_blocks); + break; + case paddle::DataType::FLOAT16: + LaunchSelectBlocks(query, + compressed_k, + compressed_k2, + block_tables, + seq_lens_decoder, + batch_id_per_token, + cu_seqlens_q, + topk_indices, + block_scores, + selected_counts, + coarse_lse, + coarse_partial_max, + coarse_partial_sum, + block_size, + kernel_size, + kernel_stride, + topk, + dense_len, + init_blocks, + local_blocks); + break; + case paddle::DataType::BFLOAT16: + LaunchSelectBlocks(query, + compressed_k, + compressed_k2, + block_tables, + seq_lens_decoder, + batch_id_per_token, + cu_seqlens_q, + topk_indices, + block_scores, + selected_counts, + coarse_lse, + coarse_partial_max, + coarse_partial_sum, + block_size, + kernel_size, + kernel_stride, + topk, + dense_len, + init_blocks, + local_blocks); + break; + default: + PD_THROW("InfLLM-V2 Stage 1 only supports fp32, fp16, or bf16."); + } + return {topk_indices, + block_scores, + selected_counts, + coarse_lse, + coarse_partial_max, + coarse_partial_sum}; +} + +std::vector InfLLMV2AttentionForward( + const paddle::Tensor& query, + const paddle::Tensor& key_cache, + const paddle::Tensor& value_cache, + const paddle::Tensor& block_tables, + const paddle::Tensor& seq_lens_decoder, + const paddle::Tensor& seq_lens_this_time, + const paddle::Tensor& batch_id_per_token, + const paddle::Tensor& cu_seqlens_q, + const paddle::Tensor& topk_indices, + paddle::Tensor& out, + paddle::Tensor& partial_acc, + paddle::Tensor& partial_max, + paddle::Tensor& partial_sum) { + for (const paddle::Tensor* tensor : + std::vector{&query, + &key_cache, + &value_cache, + &topk_indices, + &out, + &partial_acc, + &partial_max, + &partial_sum}) { + CheckCUDA(*tensor, "Stage 2 tensor"); + } + CheckFloatingType(query, "query"); + PD_CHECK(query.shape().size() == 3, + "InfLLM-V2 query must have shape [tokens, heads, head_dim]."); + PD_CHECK(key_cache.shape().size() == 4 && value_cache.shape().size() == 4, + "InfLLM-V2 K/V cache must be paged rank-4 tensors."); + PD_CHECK(query.dtype() == key_cache.dtype() && + query.dtype() == value_cache.dtype() && + query.dtype() == out.dtype(), + "InfLLM-V2 query, K/V cache, and output dtypes must match."); + PD_CHECK(key_cache.shape() == value_cache.shape(), + "InfLLM-V2 key_cache and value_cache shapes must match."); + const int64_t tokens = query.shape()[0]; + const int64_t query_heads = query.shape()[1]; + const int64_t head_dim = query.shape()[2]; + const int64_t kv_heads = key_cache.shape()[1]; + PD_CHECK(kv_heads > 0 && query_heads % kv_heads == 0, + "InfLLM-V2 query heads must be divisible by KV heads."); + PD_CHECK(head_dim == key_cache.shape()[3] && head_dim <= 256, + "InfLLM-V2 Stage 2 requires matching head_dim in [1, 256]."); + PD_CHECK(topk_indices.dtype() == paddle::DataType::INT32 && + topk_indices.shape().size() == 3 && + topk_indices.shape()[0] == tokens && + topk_indices.shape()[1] == kv_heads && + topk_indices.shape()[2] > 0, + "InfLLM-V2 topk_indices must have shape " + "[tokens, kv_heads, capacity] and dtype int32."); + PD_CHECK(out.shape() == query.shape(), + "InfLLM-V2 output workspace shape must match query."); + PD_CHECK(partial_acc.dtype() == paddle::DataType::FLOAT32 && + partial_max.dtype() == paddle::DataType::FLOAT32 && + partial_sum.dtype() == paddle::DataType::FLOAT32, + "InfLLM-V2 Stage 2 partial workspaces must be float32."); + PD_CHECK( + partial_acc.shape().size() == 4 && partial_acc.shape()[0] == tokens && + partial_acc.shape()[1] == query_heads && partial_acc.shape()[2] > 0 && + partial_acc.shape()[2] <= topk_indices.shape()[2] && + partial_acc.shape()[3] == head_dim, + "InfLLM-V2 partial_acc shape is invalid."); + PD_CHECK(partial_max.shape().size() == 3 && + partial_max.shape() == partial_sum.shape() && + partial_max.shape()[0] == tokens && + partial_max.shape()[1] == query_heads && + partial_max.shape()[2] == partial_acc.shape()[2], + "InfLLM-V2 partial max/sum shapes are invalid."); + CheckCommonMetadata(block_tables, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + static_cast(tokens)); + + switch (query.dtype()) { + case paddle::DataType::FLOAT32: + LaunchSparseAttention(query, + key_cache, + value_cache, + block_tables, + seq_lens_decoder, + batch_id_per_token, + cu_seqlens_q, + topk_indices, + out, + partial_acc, + partial_max, + partial_sum); + break; + case paddle::DataType::FLOAT16: + LaunchSparseAttention(query, + key_cache, + value_cache, + block_tables, + seq_lens_decoder, + batch_id_per_token, + cu_seqlens_q, + topk_indices, + out, + partial_acc, + partial_max, + partial_sum); + break; + case paddle::DataType::BFLOAT16: + LaunchSparseAttention(query, + key_cache, + value_cache, + block_tables, + seq_lens_decoder, + batch_id_per_token, + cu_seqlens_q, + topk_indices, + out, + partial_acc, + partial_max, + partial_sum); + break; + default: + PD_THROW("InfLLM-V2 Stage 2 only supports fp32, fp16, or bf16."); + } + return {out, partial_acc, partial_max, partial_sum}; +} + +std::vector> InfLLMV2UpdateCompressedKInferShape( + const std::vector& current_tokens_shape, + const std::vector& key_cache_shape, + const std::vector& compressed_k_shape, + const std::vector& compressed_k2_shape, + const std::vector& block_tables_shape, + const std::vector& seq_lens_decoder_shape, + const std::vector& seq_lens_this_time_shape, + const std::vector& batch_id_per_token_shape, + const std::vector& cu_seqlens_q_shape, + int kernel_size, + int kernel_stride) { + return {compressed_k_shape, compressed_k2_shape}; +} + +std::vector InfLLMV2UpdateCompressedKInferDtype( + const paddle::DataType& current_tokens_dtype, + const paddle::DataType& key_cache_dtype, + const paddle::DataType& compressed_k_dtype, + const paddle::DataType& compressed_k2_dtype, + const paddle::DataType& block_tables_dtype, + const paddle::DataType& seq_lens_decoder_dtype, + const paddle::DataType& seq_lens_this_time_dtype, + const paddle::DataType& batch_id_per_token_dtype, + const paddle::DataType& cu_seqlens_q_dtype) { + return {compressed_k_dtype, compressed_k2_dtype}; +} + +std::vector> InfLLMV2SelectBlocksInferShape( + const std::vector& query_shape, + const std::vector& compressed_k_shape, + const std::vector& compressed_k2_shape, + const std::vector& block_tables_shape, + const std::vector& seq_lens_decoder_shape, + const std::vector& seq_lens_this_time_shape, + const std::vector& batch_id_per_token_shape, + const std::vector& cu_seqlens_q_shape, + const std::vector& topk_indices_shape, + const std::vector& block_scores_shape, + const std::vector& selected_counts_shape, + const std::vector& coarse_lse_shape, + const std::vector& coarse_partial_max_shape, + const std::vector& coarse_partial_sum_shape, + int block_size, + int kernel_size, + int kernel_stride, + int topk, + int dense_len, + int init_blocks, + int local_blocks) { + return {topk_indices_shape, + block_scores_shape, + selected_counts_shape, + coarse_lse_shape, + coarse_partial_max_shape, + coarse_partial_sum_shape}; +} + +std::vector InfLLMV2SelectBlocksInferDtype( + const paddle::DataType& query_dtype, + const paddle::DataType& compressed_k_dtype, + const paddle::DataType& compressed_k2_dtype, + const paddle::DataType& block_tables_dtype, + const paddle::DataType& seq_lens_decoder_dtype, + const paddle::DataType& seq_lens_this_time_dtype, + const paddle::DataType& batch_id_per_token_dtype, + const paddle::DataType& cu_seqlens_q_dtype, + const paddle::DataType& topk_indices_dtype, + const paddle::DataType& block_scores_dtype, + const paddle::DataType& selected_counts_dtype, + const paddle::DataType& coarse_lse_dtype, + const paddle::DataType& coarse_partial_max_dtype, + const paddle::DataType& coarse_partial_sum_dtype) { + return {topk_indices_dtype, + block_scores_dtype, + selected_counts_dtype, + coarse_lse_dtype, + coarse_partial_max_dtype, + coarse_partial_sum_dtype}; +} + +std::vector> InfLLMV2AttentionForwardInferShape( + const std::vector& query_shape, + const std::vector& key_cache_shape, + const std::vector& value_cache_shape, + const std::vector& block_tables_shape, + const std::vector& seq_lens_decoder_shape, + const std::vector& seq_lens_this_time_shape, + const std::vector& batch_id_per_token_shape, + const std::vector& cu_seqlens_q_shape, + const std::vector& topk_indices_shape, + const std::vector& out_shape, + const std::vector& partial_acc_shape, + const std::vector& partial_max_shape, + const std::vector& partial_sum_shape) { + return {out_shape, partial_acc_shape, partial_max_shape, partial_sum_shape}; +} + +std::vector InfLLMV2AttentionForwardInferDtype( + const paddle::DataType& query_dtype, + const paddle::DataType& key_cache_dtype, + const paddle::DataType& value_cache_dtype, + const paddle::DataType& block_tables_dtype, + const paddle::DataType& seq_lens_decoder_dtype, + const paddle::DataType& seq_lens_this_time_dtype, + const paddle::DataType& batch_id_per_token_dtype, + const paddle::DataType& cu_seqlens_q_dtype, + const paddle::DataType& topk_indices_dtype, + const paddle::DataType& out_dtype, + const paddle::DataType& partial_acc_dtype, + const paddle::DataType& partial_max_dtype, + const paddle::DataType& partial_sum_dtype) { + return {out_dtype, partial_acc_dtype, partial_max_dtype, partial_sum_dtype}; +} + +PD_BUILD_STATIC_OP(infllmv2_update_compressed_k) + .Inputs({"current_tokens", + "key_cache", + "compressed_k", + "compressed_k2", + "block_tables", + "seq_lens_decoder", + "seq_lens_this_time", + "batch_id_per_token", + "cu_seqlens_q"}) + .Outputs({"compressed_k_out", "compressed_k2_out"}) + .SetInplaceMap({{"compressed_k", "compressed_k_out"}, + {"compressed_k2", "compressed_k2_out"}}) + .Attrs({"kernel_size: int", "kernel_stride: int"}) + .SetKernelFn(PD_KERNEL(InfLLMV2UpdateCompressedK)) + .SetInferShapeFn(PD_INFER_SHAPE(InfLLMV2UpdateCompressedKInferShape)) + .SetInferDtypeFn(PD_INFER_DTYPE(InfLLMV2UpdateCompressedKInferDtype)); + +PD_BUILD_STATIC_OP(infllmv2_select_blocks) + .Inputs({"query", + "compressed_k", + "compressed_k2", + "block_tables", + "seq_lens_decoder", + "seq_lens_this_time", + "batch_id_per_token", + "cu_seqlens_q", + "topk_indices", + "block_scores", + "selected_counts", + "coarse_lse", + "coarse_partial_max", + "coarse_partial_sum"}) + .Outputs({"topk_indices_out", + "block_scores_out", + "selected_counts_out", + "coarse_lse_out", + "coarse_partial_max_out", + "coarse_partial_sum_out"}) + .SetInplaceMap({{"topk_indices", "topk_indices_out"}, + {"block_scores", "block_scores_out"}, + {"selected_counts", "selected_counts_out"}, + {"coarse_lse", "coarse_lse_out"}, + {"coarse_partial_max", "coarse_partial_max_out"}, + {"coarse_partial_sum", "coarse_partial_sum_out"}}) + .Attrs({"block_size: int", + "kernel_size: int", + "kernel_stride: int", + "topk: int", + "dense_len: int", + "init_blocks: int", + "local_blocks: int"}) + .SetKernelFn(PD_KERNEL(InfLLMV2SelectBlocks)) + .SetInferShapeFn(PD_INFER_SHAPE(InfLLMV2SelectBlocksInferShape)) + .SetInferDtypeFn(PD_INFER_DTYPE(InfLLMV2SelectBlocksInferDtype)); + +PD_BUILD_STATIC_OP(infllmv2_attention_forward) + .Inputs({"query", + "key_cache", + "value_cache", + "block_tables", + "seq_lens_decoder", + "seq_lens_this_time", + "batch_id_per_token", + "cu_seqlens_q", + "topk_indices", + "out", + "partial_acc", + "partial_max", + "partial_sum"}) + .Outputs( + {"out_alias", "partial_acc_out", "partial_max_out", "partial_sum_out"}) + .SetInplaceMap({{"out", "out_alias"}, + {"partial_acc", "partial_acc_out"}, + {"partial_max", "partial_max_out"}, + {"partial_sum", "partial_sum_out"}}) + .SetKernelFn(PD_KERNEL(InfLLMV2AttentionForward)) + .SetInferShapeFn(PD_INFER_SHAPE(InfLLMV2AttentionForwardInferShape)) + .SetInferDtypeFn(PD_INFER_DTYPE(InfLLMV2AttentionForwardInferDtype)); diff --git a/custom_ops/gpu_ops/infllmv2_attention/infllmv2_impl.cuh b/custom_ops/gpu_ops/infllmv2_attention/infllmv2_impl.cuh new file mode 100644 index 00000000000..2c28c0deb04 --- /dev/null +++ b/custom_ops/gpu_ops/infllmv2_attention/infllmv2_impl.cuh @@ -0,0 +1,1431 @@ +// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace fastdeploy { + +template +__device__ __forceinline__ float InfLLMToFloat(T value) { + return static_cast(value); +} + +template <> +__device__ __forceinline__ float InfLLMToFloat(half value) { + return __half2float(value); +} + +template <> +__device__ __forceinline__ float InfLLMToFloat<__nv_bfloat16>( + __nv_bfloat16 value) { + return __bfloat162float(value); +} + +template +__device__ __forceinline__ T InfLLMFromFloat(float value) { + return static_cast(value); +} + +template <> +__device__ __forceinline__ half InfLLMFromFloat(float value) { + return __float2half_rn(value); +} + +template <> +__device__ __forceinline__ __nv_bfloat16 +InfLLMFromFloat<__nv_bfloat16>(float value) { + return __float2bfloat16_rn(value); +} + +__device__ __forceinline__ int InfLLMQueryPosition(int token_id, + int batch_id, + const int* seq_lens_decoder, + const int* cu_seqlens_q) { + return seq_lens_decoder[batch_id] + token_id - cu_seqlens_q[batch_id]; +} + +template +__device__ __forceinline__ float InfLLMReadPaged(const T* cache, + const int* block_tables, + int max_blocks_per_seq, + int batch_id, + int logical_position, + int kv_head, + int kv_heads, + int block_size, + int head_dim, + int dim) { + const int logical_block = logical_position / block_size; + const int block_offset = logical_position % block_size; + const int physical_block = + block_tables[batch_id * max_blocks_per_seq + logical_block]; + const int64_t offset = + ((static_cast(physical_block) * kv_heads + kv_head) * + block_size + + block_offset) * + head_dim + + dim; + return InfLLMToFloat(cache[offset]); +} + +template +__global__ void InfLLMV2UpdateCompressedKKernel(const T* key_cache, + T* compressed_k, + T* compressed_k2, + const int* block_tables, + const int* seq_lens_decoder, + const int* batch_id_per_token, + const int* cu_seqlens_q, + int tokens, + int batch_size, + int physical_blocks, + int max_blocks_per_seq, + int kv_heads, + int block_size, + int head_dim, + int kernel_size, + int kernel_stride) { + const int token_id = blockIdx.x; + const int kv_head = blockIdx.y; + const int scale_id = blockIdx.z; + const int dim = threadIdx.x; + if (token_id >= tokens || kv_head >= kv_heads || dim >= head_dim) { + return; + } + + const int batch_id = batch_id_per_token[token_id]; + if (batch_id < 0 || batch_id >= batch_size) { + return; + } + const int position = + InfLLMQueryPosition(token_id, batch_id, seq_lens_decoder, cu_seqlens_q); + const int window_size = scale_id == 0 ? kernel_size : 4 * kernel_size; + const int stride = scale_id == 0 ? kernel_stride : 4 * kernel_stride; + if (position + 1 < window_size || + (position + 1 - window_size) % stride != 0) { + return; + } + + float sum = 0.0f; + const int window_begin = position + 1 - window_size; + for (int logical_position = window_begin; logical_position <= position; + ++logical_position) { + sum += InfLLMReadPaged(key_cache, + block_tables, + max_blocks_per_seq, + batch_id, + logical_position, + kv_head, + kv_heads, + block_size, + head_dim, + dim); + } + + const int logical_block = position / block_size; + const int physical_block = + block_tables[batch_id * max_blocks_per_seq + logical_block]; + if (physical_block < 0 || physical_block >= physical_blocks) { + return; + } + const int slots = block_size / stride; + const int slot = (position / stride) % slots; + const int64_t output_offset = + ((static_cast(physical_block) * kv_heads + kv_head) * slots + + slot) * + head_dim + + dim; + T* output = scale_id == 0 ? compressed_k : compressed_k2; + output[output_offset] = InfLLMFromFloat(sum / window_size); +} + +template +__device__ __forceinline__ float InfLLMReadSummary(const T* compressed, + const int* block_tables, + int max_blocks_per_seq, + int batch_id, + int window_end, + int stride, + int slots, + int kv_head, + int kv_heads, + int head_dim, + int dim) { + const int logical_block = window_end / (slots * stride); + const int physical_block = + block_tables[batch_id * max_blocks_per_seq + logical_block]; + const int slot = (window_end / stride) % slots; + const int64_t offset = + ((static_cast(physical_block) * kv_heads + kv_head) * slots + + slot) * + head_dim + + dim; + return InfLLMToFloat(compressed[offset]); +} + +template +__global__ void InfLLMV2CoarseLSETensorCoreSplitKernel( + const T* query, + const T* compressed_k2, + const int* block_tables, + const int* seq_lens_decoder, + const int* batch_id_per_token, + const int* cu_seqlens_q, + float* coarse_partial_max, + float* coarse_partial_sum, + int tokens, + int batch_size, + int max_blocks_per_seq, + int query_heads, + int kv_heads, + int block_size, + int head_dim, + int kernel_size, + int kernel_stride, + int coarse_splits) { + constexpr int kTensorTile = 16; + const int token_id = blockIdx.x; + const int kv_head = blockIdx.y; + const int split_id = blockIdx.z; + const int lane = threadIdx.x; + if (token_id >= tokens || kv_head >= kv_heads || split_id >= coarse_splits) { + return; + } + const int batch_id = batch_id_per_token[token_id]; + if (batch_id < 0 || batch_id >= batch_size) { + if (lane < kTensorTile) { + const int query_head = kv_head * kTensorTile + lane; + const int partial_offset = + (token_id * query_heads + query_head) * coarse_splits + split_id; + coarse_partial_max[partial_offset] = -CUDART_INF_F; + coarse_partial_sum[partial_offset] = 0.0f; + } + return; + } + + const int position = + InfLLMQueryPosition(token_id, batch_id, seq_lens_decoder, cu_seqlens_q); + const int visible_length = position + 1; + const int coarse_kernel = 4 * kernel_size; + const int coarse_stride = 4 * kernel_stride; + const int coarse_slots = block_size / coarse_stride; + const int windows = + visible_length < coarse_kernel + ? 0 + : (visible_length - coarse_kernel) / coarse_stride + 1; + const int first_window = split_id * kTensorTile; + const int window_count = max(0, min(kTensorTile, windows - first_window)); + extern __shared__ char coarse_tensor_shared[]; + T* shared_k = reinterpret_cast(coarse_tensor_shared); + const size_t scores_offset = + (static_cast(kTensorTile) * head_dim * sizeof(T) + + alignof(float) - 1) & + ~(alignof(float) - 1); + float* shared_scores = + reinterpret_cast(coarse_tensor_shared + scores_offset); + + for (int index = lane; index < kTensorTile * head_dim; index += 32) { + const int local_window = index / head_dim; + const int dim = index % head_dim; + if (local_window < window_count) { + const int window = first_window + local_window; + const int window_end = coarse_kernel - 1 + window * coarse_stride; + shared_k[index] = InfLLMReadSummary(compressed_k2, + block_tables, + max_blocks_per_seq, + batch_id, + window_end, + coarse_stride, + coarse_slots, + kv_head, + kv_heads, + head_dim, + dim); + } else { + shared_k[index] = InfLLMFromFloat(0.0f); + } + } + __syncwarp(); + + using namespace nvcuda; + wmma:: + fragment + accumulator; + wmma::fill_fragment(accumulator, 0.0f); + const T* query_base = query + (static_cast(token_id) * query_heads + + kv_head * kTensorTile) * + head_dim; + for (int dim = 0; dim < head_dim; dim += kTensorTile) { + wmma::fragment + query_fragment; + wmma::fragment + key_fragment; + wmma::load_matrix_sync(query_fragment, query_base + dim, head_dim); + wmma::load_matrix_sync(key_fragment, shared_k + dim, head_dim); + wmma::mma_sync(accumulator, query_fragment, key_fragment, accumulator); + } + wmma::store_matrix_sync( + shared_scores, accumulator, kTensorTile, wmma::mem_row_major); + __syncwarp(); + + if (lane < kTensorTile) { + const int query_head = kv_head * kTensorTile + lane; + const int score_base = lane * kTensorTile; + const float scale = rsqrtf(static_cast(head_dim)); + float row_max = -CUDART_INF_F; + for (int local_window = 0; local_window < window_count; ++local_window) { + row_max = + fmaxf(row_max, shared_scores[score_base + local_window] * scale); + } + float row_sum = 0.0f; + for (int local_window = 0; local_window < window_count; ++local_window) { + row_sum += + expf(shared_scores[score_base + local_window] * scale - row_max); + } + const int partial_offset = + (token_id * query_heads + query_head) * coarse_splits + split_id; + coarse_partial_max[partial_offset] = row_max; + coarse_partial_sum[partial_offset] = row_sum; + } +} + +__global__ void InfLLMV2CoarseLSECombineKernel(const float* coarse_partial_max, + const float* coarse_partial_sum, + const int* batch_id_per_token, + float* coarse_lse, + int tokens, + int batch_size, + int query_heads, + int coarse_splits) { + const int token_id = blockIdx.x; + const int query_head = blockIdx.y; + const int thread = threadIdx.x; + if (token_id >= tokens || query_head >= query_heads) { + return; + } + const int partial_base = + (token_id * query_heads + query_head) * coarse_splits; + float local_max = -CUDART_INF_F; + for (int split = thread; split < coarse_splits; split += blockDim.x) { + local_max = fmaxf(local_max, coarse_partial_max[partial_base + split]); + } + extern __shared__ float reduction[]; + float* maxima = reduction; + float* sums = reduction + blockDim.x; + maxima[thread] = local_max; + __syncthreads(); + for (int width = blockDim.x / 2; width > 0; width >>= 1) { + if (thread < width) { + maxima[thread] = fmaxf(maxima[thread], maxima[thread + width]); + } + __syncthreads(); + } + const float row_max = maxima[0]; + float local_sum = 0.0f; + if (row_max != -CUDART_INF_F) { + for (int split = thread; split < coarse_splits; split += blockDim.x) { + const float split_sum = coarse_partial_sum[partial_base + split]; + if (split_sum > 0.0f) { + local_sum += split_sum * + expf(coarse_partial_max[partial_base + split] - row_max); + } + } + } + sums[thread] = local_sum; + __syncthreads(); + for (int width = blockDim.x / 2; width > 0; width >>= 1) { + if (thread < width) { + sums[thread] += sums[thread + width]; + } + __syncthreads(); + } + if (thread == 0) { + const int batch_id = batch_id_per_token[token_id]; + coarse_lse[token_id * query_heads + query_head] = + batch_id < 0 || batch_id >= batch_size + ? -CUDART_INF_F + : (row_max == -CUDART_INF_F ? 0.0f : row_max + logf(sums[0])); + } +} + +template +__global__ void InfLLMV2CoarseLSEKernel(const T* query, + const T* compressed_k2, + const int* block_tables, + const int* seq_lens_decoder, + const int* batch_id_per_token, + const int* cu_seqlens_q, + float* coarse_lse, + float* coarse_partial_max, + float* coarse_partial_sum, + int tokens, + int batch_size, + int max_blocks_per_seq, + int query_heads, + int kv_heads, + int block_size, + int head_dim, + int kernel_size, + int kernel_stride, + int coarse_splits) { + const int token_id = blockIdx.x; + const int query_head = blockIdx.y; + if (token_id >= tokens || query_head >= query_heads) { + return; + } + const int batch_id = batch_id_per_token[token_id]; + const int partial_base = + (token_id * query_heads + query_head) * coarse_splits; + for (int split = threadIdx.x; split < coarse_splits; split += blockDim.x) { + coarse_partial_max[partial_base + split] = -CUDART_INF_F; + coarse_partial_sum[partial_base + split] = 0.0f; + } + if (batch_id < 0 || batch_id >= batch_size) { + if (threadIdx.x == 0) { + coarse_lse[token_id * query_heads + query_head] = -CUDART_INF_F; + } + return; + } + + const int position = + InfLLMQueryPosition(token_id, batch_id, seq_lens_decoder, cu_seqlens_q); + const int visible_length = position + 1; + const int coarse_kernel = 4 * kernel_size; + const int coarse_stride = 4 * kernel_stride; + const int coarse_slots = block_size / coarse_stride; + const int windows = + visible_length < coarse_kernel + ? 0 + : (visible_length - coarse_kernel) / coarse_stride + 1; + const int kv_head = query_head / (query_heads / kv_heads); + const float scale = rsqrtf(static_cast(head_dim)); + + float local_max = -CUDART_INF_F; + float local_sum = 0.0f; + for (int window = threadIdx.x; window < windows; window += blockDim.x) { + const int window_end = coarse_kernel - 1 + window * coarse_stride; + float dot = 0.0f; + for (int dim = 0; dim < head_dim; ++dim) { + const int64_t query_offset = + (static_cast(token_id) * query_heads + query_head) * + head_dim + + dim; + dot += InfLLMToFloat(query[query_offset]) * + InfLLMReadSummary(compressed_k2, + block_tables, + max_blocks_per_seq, + batch_id, + window_end, + coarse_stride, + coarse_slots, + kv_head, + kv_heads, + head_dim, + dim); + } + const float score = dot * scale; + if (score > local_max) { + local_sum = local_max == -CUDART_INF_F + ? 1.0f + : local_sum * expf(local_max - score) + 1.0f; + local_max = score; + } else { + local_sum += expf(score - local_max); + } + } + + extern __shared__ float reduction[]; + float* maxima = reduction; + float* sums = reduction + blockDim.x; + maxima[threadIdx.x] = local_max; + sums[threadIdx.x] = local_sum; + __syncthreads(); + for (int width = blockDim.x / 2; width > 0; width >>= 1) { + if (threadIdx.x < width) { + const float other_max = maxima[threadIdx.x + width]; + const float merged_max = fmaxf(maxima[threadIdx.x], other_max); + float merged_sum = 0.0f; + if (merged_max != -CUDART_INF_F) { + if (maxima[threadIdx.x] != -CUDART_INF_F) { + merged_sum += + sums[threadIdx.x] * expf(maxima[threadIdx.x] - merged_max); + } + if (other_max != -CUDART_INF_F) { + merged_sum += + sums[threadIdx.x + width] * expf(other_max - merged_max); + } + } + maxima[threadIdx.x] = merged_max; + sums[threadIdx.x] = merged_sum; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + coarse_lse[token_id * query_heads + query_head] = + windows == 0 ? 0.0f : maxima[0] + logf(sums[0]); + } +} + +template +__global__ void InfLLMV2BlockScoreTensorCoreKernel( + const T* query, + const T* compressed_k, + const int* block_tables, + const int* seq_lens_decoder, + const int* batch_id_per_token, + const int* cu_seqlens_q, + const float* coarse_lse, + float* block_scores, + int tokens, + int batch_size, + int max_blocks_per_seq, + int query_heads, + int kv_heads, + int block_size, + int head_dim, + int kernel_size, + int kernel_stride, + int init_blocks, + int local_blocks) { + constexpr int kWarpsPerCTA = 4; + constexpr int kTensorTile = 16; + const int token_id = blockIdx.x; + const int kv_head = blockIdx.y; + const int warp_id = threadIdx.x / 32; + const int lane = threadIdx.x % 32; + const int logical_block = blockIdx.z * kWarpsPerCTA + warp_id; + if (token_id >= tokens || kv_head >= kv_heads || + logical_block >= max_blocks_per_seq) { + return; + } + const int output_offset = + (token_id * kv_heads + kv_head) * max_blocks_per_seq + logical_block; + const int batch_id = batch_id_per_token[token_id]; + if (batch_id < 0 || batch_id >= batch_size) { + if (lane == 0) { + block_scores[output_offset] = -CUDART_INF_F; + } + return; + } + const int position = + InfLLMQueryPosition(token_id, batch_id, seq_lens_decoder, cu_seqlens_q); + const int visible_length = position + 1; + const int valid_blocks = (visible_length + block_size - 1) / block_size; + if (logical_block >= valid_blocks) { + if (lane == 0) { + block_scores[output_offset] = -CUDART_INF_F; + } + return; + } + const int current_block = position / block_size; + if (logical_block < init_blocks || + (logical_block <= current_block && + logical_block + local_blocks > current_block)) { + if (lane == 0) { + block_scores[output_offset] = CUDART_INF_F; + } + return; + } + + const int fine_slots = block_size / kernel_stride; + const int first_window = max(0, logical_block * fine_slots - 1); + const int fine_windows = + visible_length < kernel_size + ? 0 + : (visible_length - kernel_size) / kernel_stride + 1; + const int last_window = min(fine_windows, (logical_block + 1) * fine_slots); + const int window_count = last_window - first_window; + extern __shared__ char tensor_score_shared[]; + const size_t shared_k_bytes = + static_cast(kTensorTile) * head_dim * sizeof(T); + const size_t scores_offset = + (shared_k_bytes + alignof(float) - 1) & ~(alignof(float) - 1); + const size_t warp_bytes = + scores_offset + kTensorTile * kTensorTile * sizeof(float); + char* warp_shared = tensor_score_shared + warp_id * warp_bytes; + T* shared_k = reinterpret_cast(warp_shared); + float* shared_scores = reinterpret_cast(warp_shared + scores_offset); + + for (int index = lane; index < kTensorTile * head_dim; index += 32) { + const int local_window = index / head_dim; + const int dim = index % head_dim; + if (local_window < window_count) { + const int window = first_window + local_window; + const int window_end = kernel_size - 1 + window * kernel_stride; + shared_k[index] = InfLLMReadSummary(compressed_k, + block_tables, + max_blocks_per_seq, + batch_id, + window_end, + kernel_stride, + fine_slots, + kv_head, + kv_heads, + head_dim, + dim); + } else { + shared_k[index] = InfLLMFromFloat(0.0f); + } + } + __syncwarp(); + + using namespace nvcuda; + wmma:: + fragment + accumulator; + wmma::fill_fragment(accumulator, 0.0f); + const int query_head = kv_head * kTensorTile; + const T* query_base = + query + + (static_cast(token_id) * query_heads + query_head) * head_dim; + for (int dim = 0; dim < head_dim; dim += kTensorTile) { + wmma::fragment + query_fragment; + wmma::fragment + key_fragment; + wmma::load_matrix_sync(query_fragment, query_base + dim, head_dim); + wmma::load_matrix_sync(key_fragment, shared_k + dim, head_dim); + wmma::mma_sync(accumulator, query_fragment, key_fragment, accumulator); + } + wmma::store_matrix_sync( + shared_scores, accumulator, kTensorTile, wmma::mem_row_major); + __syncwarp(); + + float gqa_score = -CUDART_INF_F; + if (lane < window_count) { + gqa_score = 0.0f; + const float scale = rsqrtf(static_cast(head_dim)); + for (int head = 0; head < kTensorTile; ++head) { + const float lse = coarse_lse[token_id * query_heads + query_head + head]; + gqa_score += expf(shared_scores[head * kTensorTile + lane] * scale - lse); + } + } + for (int delta = 16; delta > 0; delta >>= 1) { + gqa_score = + fmaxf(gqa_score, __shfl_down_sync(0xffffffff, gqa_score, delta)); + } + if (lane == 0) { + block_scores[output_offset] = gqa_score; + } +} + +template +__global__ void InfLLMV2BlockScoreKernel(const T* query, + const T* compressed_k, + const int* block_tables, + const int* seq_lens_decoder, + const int* batch_id_per_token, + const int* cu_seqlens_q, + const float* coarse_lse, + float* block_scores, + int tokens, + int batch_size, + int max_blocks_per_seq, + int query_heads, + int kv_heads, + int block_size, + int head_dim, + int kernel_size, + int kernel_stride, + int init_blocks, + int local_blocks) { + extern __shared__ char block_score_shared[]; + constexpr int kBlocksPerCTA = 4; + const int token_id = blockIdx.x; + const int kv_head = blockIdx.y; + const int first_logical_block = blockIdx.z * kBlocksPerCTA; + if (token_id >= tokens || kv_head >= kv_heads || + first_logical_block >= max_blocks_per_seq) { + return; + } + const int output_base = (token_id * kv_heads + kv_head) * max_blocks_per_seq; + const int batch_id = batch_id_per_token[token_id]; + if (batch_id < 0 || batch_id >= batch_size) { + if (threadIdx.x == 0) { + for (int offset = 0; offset < kBlocksPerCTA; ++offset) { + const int logical_block = first_logical_block + offset; + if (logical_block < max_blocks_per_seq) { + block_scores[output_base + logical_block] = -CUDART_INF_F; + } + } + } + return; + } + const int position = + InfLLMQueryPosition(token_id, batch_id, seq_lens_decoder, cu_seqlens_q); + const int visible_length = position + 1; + const int valid_blocks = (visible_length + block_size - 1) / block_size; + const int current_block = position / block_size; + const int fine_slots = block_size / kernel_stride; + const int fine_windows = + visible_length < kernel_size + ? 0 + : (visible_length - kernel_size) / kernel_stride + 1; + const int group_size = query_heads / kv_heads; + const int group_head = threadIdx.x / 32; + const int lane = threadIdx.x % 32; + const int query_head = kv_head * group_size + group_head; + const int values_per_lane = (head_dim + 31) / 32; + float query_values[8]; +#pragma unroll + for (int element = 0; element < 8; ++element) { + query_values[element] = 0.0f; + } + const int64_t query_offset = + (static_cast(token_id) * query_heads + query_head) * head_dim; + for (int element = 0; element < values_per_lane; ++element) { + const int dim = lane + element * 32; + if (dim < head_dim) { + query_values[element] = InfLLMToFloat(query[query_offset + dim]); + } + } + T* shared_k = reinterpret_cast(block_score_shared); + const size_t shared_k_bytes = + static_cast(fine_slots + 1) * head_dim * sizeof(T); + const size_t scores_offset = + (shared_k_bytes + alignof(float) - 1) & ~(alignof(float) - 1); + float* gqa_scores = + reinterpret_cast(block_score_shared + scores_offset); + const float scale = rsqrtf(static_cast(head_dim)); + const float lse = coarse_lse[token_id * query_heads + query_head]; + + for (int block_offset = 0; block_offset < kBlocksPerCTA; ++block_offset) { + const int logical_block = first_logical_block + block_offset; + if (logical_block >= max_blocks_per_seq) { + break; + } + const int output_offset = output_base + logical_block; + if (logical_block >= valid_blocks) { + if (threadIdx.x == 0) { + block_scores[output_offset] = -CUDART_INF_F; + } + continue; + } + if (logical_block < init_blocks || + (logical_block <= current_block && + logical_block + local_blocks > current_block)) { + if (threadIdx.x == 0) { + block_scores[output_offset] = CUDART_INF_F; + } + continue; + } + + const int first_window = max(0, logical_block * fine_slots - 1); + const int last_window = min(fine_windows, (logical_block + 1) * fine_slots); + const int window_count = last_window - first_window; + for (int index = threadIdx.x; index < window_count * head_dim; + index += blockDim.x) { + const int local_window = index / head_dim; + const int dim = index % head_dim; + const int window = first_window + local_window; + const int window_end = kernel_size - 1 + window * kernel_stride; + shared_k[index] = InfLLMReadSummary(compressed_k, + block_tables, + max_blocks_per_seq, + batch_id, + window_end, + kernel_stride, + fine_slots, + kv_head, + kv_heads, + head_dim, + dim); + } + __syncthreads(); + + for (int local_window = 0; local_window < window_count; ++local_window) { + float local_dot = 0.0f; + for (int element = 0; element < values_per_lane; ++element) { + const int dim = lane + element * 32; + if (dim < head_dim) { + local_dot += query_values[element] * + InfLLMToFloat(shared_k[local_window * head_dim + dim]); + } + } + for (int delta = 16; delta > 0; delta >>= 1) { + local_dot += __shfl_down_sync(0xffffffff, local_dot, delta); + } + if (lane == 0) { + gqa_scores[local_window * group_size + group_head] = + expf(local_dot * scale - lse); + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + float best = -CUDART_INF_F; + for (int local_window = 0; local_window < window_count; ++local_window) { + float gqa_score = 0.0f; + for (int head = 0; head < group_size; ++head) { + gqa_score += gqa_scores[local_window * group_size + head]; + } + best = fmaxf(best, gqa_score); + } + block_scores[output_offset] = best; + } + __syncthreads(); + } +} + +__device__ __forceinline__ uint32_t InfLLMOrderedFloatBits(float value) { + const uint32_t bits = __float_as_uint(value); + return bits ^ ((static_cast(bits) < 0) ? 0xffffffffu : 0x80000000u); +} + +__global__ void InfLLMV2TopKKernel(const float* block_scores, + const int* seq_lens_decoder, + const int* batch_id_per_token, + const int* cu_seqlens_q, + int* topk_indices, + int* selected_counts, + int tokens, + int batch_size, + int kv_heads, + int max_blocks_per_seq, + int capacity, + int block_size, + int topk, + int dense_len, + int local_blocks) { + constexpr int kBlockThreads = 256; + constexpr int kItemsPerThread = 8; + using BlockRadixSort = + cub::BlockRadixSort; + __shared__ typename BlockRadixSort::TempStorage sort_storage; + + const int token_id = blockIdx.x; + const int kv_head = blockIdx.y; + if (token_id >= tokens || kv_head >= kv_heads) { + return; + } + int* selected = topk_indices + (token_id * kv_heads + kv_head) * capacity; + for (int slot = threadIdx.x; slot < capacity; slot += blockDim.x) { + selected[slot] = -1; + } + const int count_offset = token_id * kv_heads + kv_head; + const int batch_id = batch_id_per_token[token_id]; + if (batch_id < 0 || batch_id >= batch_size) { + if (threadIdx.x == 0) { + selected_counts[count_offset] = 0; + } + return; + } + const int position = + InfLLMQueryPosition(token_id, batch_id, seq_lens_decoder, cu_seqlens_q); + const int visible_length = position + 1; + const int valid_blocks = (visible_length + block_size - 1) / block_size; + const int target = + min(valid_blocks, + visible_length < dense_len ? valid_blocks : topk + local_blocks); + const float* scores = + block_scores + (token_id * kv_heads + kv_head) * max_blocks_per_seq; + + uint64_t sort_keys[kItemsPerThread]; +#pragma unroll + for (int item = 0; item < kItemsPerThread; ++item) { + const int candidate = threadIdx.x * kItemsPerThread + item; + if (candidate < valid_blocks) { + const uint64_t score_key = InfLLMOrderedFloatBits(scores[candidate]); + const uint64_t tie_key = 0xffffffffu - candidate; + sort_keys[item] = (score_key << 32) | tie_key; + } else { + sort_keys[item] = 0; + } + } + + BlockRadixSort(sort_storage).SortDescending(sort_keys); + __syncthreads(); + +#pragma unroll + for (int item = 0; item < kItemsPerThread; ++item) { + const int rank = threadIdx.x * kItemsPerThread + item; + const uint32_t candidate = + 0xffffffffu - static_cast(sort_keys[item]); + sort_keys[item] = rank < target ? candidate : UINT64_MAX; + } + + BlockRadixSort(sort_storage).Sort(sort_keys); + +#pragma unroll + for (int item = 0; item < kItemsPerThread; ++item) { + const int rank = threadIdx.x * kItemsPerThread + item; + if (rank < target) { + selected[rank] = static_cast(sort_keys[item]); + } + } + if (threadIdx.x == 0) { + selected_counts[count_offset] = target; + } +} + +template +__global__ void InfLLMV2SparseAttentionTensorCoreSplitKVKernel( + const T* query, + const T* key_cache, + const T* value_cache, + const int* block_tables, + const int* seq_lens_decoder, + const int* batch_id_per_token, + const int* cu_seqlens_q, + const int* topk_indices, + float* partial_acc, + float* partial_max, + float* partial_sum, + int tokens, + int batch_size, + int physical_blocks, + int max_blocks_per_seq, + int query_heads, + int kv_heads, + int block_size, + int head_dim, + int capacity, + int splits, + int blocks_per_split) { + constexpr int kTokenTile = 16; + constexpr int kGroupSize = 16; + constexpr int kWarpsPerCTA = 4; + constexpr int kHeadsPerWarp = kGroupSize / kWarpsPerCTA; + constexpr int kElementsPerVector = sizeof(uint4) / sizeof(T); + static_assert(sizeof(T) == 2, + "Tensor-core Stage 2 requires a 16-bit cache dtype."); + const int token_id = blockIdx.x; + const int kv_head = blockIdx.y; + const int split_id = blockIdx.z; + const int warp_id = threadIdx.x / 32; + const int lane = threadIdx.x % 32; + if (token_id >= tokens || kv_head >= kv_heads || split_id >= splits) { + return; + } + const int first_query_head = kv_head * kGroupSize + warp_id * kHeadsPerWarp; + const int values_per_lane = (head_dim + 31) / 32; + + const int batch_id = batch_id_per_token[token_id]; + if (batch_id < 0 || batch_id >= batch_size) { + for (int head = 0; head < kHeadsPerWarp; ++head) { + const int query_head = first_query_head + head; + const int64_t scalar_offset = + (static_cast(token_id) * query_heads + query_head) * splits + + split_id; + const int64_t acc_offset = scalar_offset * head_dim; + for (int element = 0; element < values_per_lane; ++element) { + const int dim = lane + element * 32; + if (dim < head_dim) { + partial_acc[acc_offset + dim] = 0.0f; + } + } + if (lane == 0) { + partial_max[scalar_offset] = -CUDART_INF_F; + partial_sum[scalar_offset] = 0.0f; + } + } + return; + } + + const int position = + InfLLMQueryPosition(token_id, batch_id, seq_lens_decoder, cu_seqlens_q); + const int* selected = + topk_indices + (token_id * kv_heads + kv_head) * capacity; + const int first_slot = split_id * blocks_per_split; + const int last_slot = min(capacity, first_slot + blocks_per_split); + const float scale = rsqrtf(static_cast(head_dim)); + extern __shared__ char tensor_attention_shared[]; + T* shared_key = reinterpret_cast(tensor_attention_shared); + T* shared_value = shared_key + kTokenTile * head_dim; + const size_t scores_offset = + (static_cast(2 * kTokenTile * head_dim) * sizeof(T) + + alignof(float) - 1) & + ~(alignof(float) - 1); + float* shared_scores = + reinterpret_cast(tensor_attention_shared + scores_offset); + const size_t probability_offset = + scores_offset + kGroupSize * kTokenTile * sizeof(float); + T* shared_probability = + reinterpret_cast(tensor_attention_shared + probability_offset); + const size_t output_offset = + probability_offset + kGroupSize * kTokenTile * sizeof(T); + float* shared_output = + reinterpret_cast(tensor_attention_shared + output_offset); + float* shared_pv = shared_output + kGroupSize * head_dim; + float* shared_row_max = shared_pv + kWarpsPerCTA * kGroupSize * kTokenTile; + float* shared_row_sum = shared_row_max + kGroupSize; + float* shared_old_scale = shared_row_sum + kGroupSize; + + for (int index = threadIdx.x; index < kGroupSize * head_dim; + index += blockDim.x) { + shared_output[index] = 0.0f; + } + if (threadIdx.x < kGroupSize) { + shared_row_max[threadIdx.x] = -CUDART_INF_F; + shared_row_sum[threadIdx.x] = 0.0f; + } + __syncthreads(); + + for (int slot = first_slot; slot < last_slot; ++slot) { + const int logical_block = selected[slot]; + if (logical_block < 0) { + break; + } + const int physical_block = + block_tables[batch_id * max_blocks_per_seq + logical_block]; + if (physical_block < 0 || physical_block >= physical_blocks) { + continue; + } + for (int block_offset = 0; block_offset < block_size; + block_offset += kTokenTile) { + const int first_position = logical_block * block_size + block_offset; + if (first_position > position) { + break; + } + const int tile_tokens = + min(kTokenTile, + min(block_size - block_offset, position - first_position + 1)); + const int64_t cache_tile_offset = + ((static_cast(physical_block) * kv_heads + kv_head) * + block_size + + block_offset) * + head_dim; + if (tile_tokens == kTokenTile) { + const int tile_vectors = kTokenTile * head_dim / kElementsPerVector; + for (int vector_index = threadIdx.x; vector_index < tile_vectors; + vector_index += blockDim.x) { + const int element_index = vector_index * kElementsPerVector; + const uint4 packed_key = *reinterpret_cast( + key_cache + cache_tile_offset + element_index); + const uint4 packed_value = *reinterpret_cast( + value_cache + cache_tile_offset + element_index); + *reinterpret_cast(shared_key + element_index) = packed_key; + *reinterpret_cast(shared_value + element_index) = + packed_value; + } + } else { + for (int index = threadIdx.x; index < kTokenTile * head_dim; + index += blockDim.x) { + const int tile_token = index / head_dim; + if (tile_token < tile_tokens) { + shared_key[index] = key_cache[cache_tile_offset + index]; + shared_value[index] = value_cache[cache_tile_offset + index]; + } else { + shared_key[index] = InfLLMFromFloat(0.0f); + shared_value[index] = InfLLMFromFloat(0.0f); + } + } + } + __syncthreads(); + + if (warp_id == 0) { + using namespace nvcuda; + wmma::fragment + score_fragment; + wmma::fill_fragment(score_fragment, 0.0f); + const T* query_base = + query + (static_cast(token_id) * query_heads + + kv_head * kGroupSize) * + head_dim; + for (int dim = 0; dim < head_dim; dim += kTokenTile) { + wmma::fragment + query_fragment; + wmma::fragment + key_fragment; + wmma::load_matrix_sync(query_fragment, query_base + dim, head_dim); + wmma::load_matrix_sync(key_fragment, shared_key + dim, head_dim); + wmma::mma_sync( + score_fragment, query_fragment, key_fragment, score_fragment); + } + wmma::store_matrix_sync( + shared_scores, score_fragment, kTokenTile, wmma::mem_row_major); + } + __syncthreads(); + + if (threadIdx.x < kGroupSize) { + const int group_head = threadIdx.x; + const int score_base = group_head * kTokenTile; + float tile_max = -CUDART_INF_F; + for (int tile_token = 0; tile_token < tile_tokens; ++tile_token) { + tile_max = + fmaxf(tile_max, shared_scores[score_base + tile_token] * scale); + } + const float previous_max = shared_row_max[group_head]; + const float next_max = fmaxf(previous_max, tile_max); + const float old_scale = previous_max == -CUDART_INF_F + ? 0.0f + : expf(previous_max - next_max); + float tile_sum = 0.0f; + for (int tile_token = 0; tile_token < kTokenTile; ++tile_token) { + const float probability = + tile_token < tile_tokens + ? expf(shared_scores[score_base + tile_token] * scale - + next_max) + : 0.0f; + shared_probability[score_base + tile_token] = + InfLLMFromFloat(probability); + tile_sum += probability; + } + shared_old_scale[group_head] = old_scale; + shared_row_max[group_head] = next_max; + shared_row_sum[group_head] = + shared_row_sum[group_head] * old_scale + tile_sum; + } + __syncthreads(); + + using namespace nvcuda; + for (int dim_tile = warp_id; dim_tile < head_dim / kTokenTile; + dim_tile += kWarpsPerCTA) { + wmma::fragment + probability_fragment; + wmma::fragment + value_fragment; + wmma::fragment + output_fragment; + wmma::fill_fragment(output_fragment, 0.0f); + wmma::load_matrix_sync( + probability_fragment, shared_probability, kTokenTile); + wmma::load_matrix_sync( + value_fragment, shared_value + dim_tile * kTokenTile, head_dim); + wmma::mma_sync(output_fragment, + probability_fragment, + value_fragment, + output_fragment); + float* warp_pv = shared_pv + warp_id * kGroupSize * kTokenTile; + wmma::store_matrix_sync( + warp_pv, output_fragment, kTokenTile, wmma::mem_row_major); + __syncwarp(); + for (int index = lane; index < kGroupSize * kTokenTile; index += 32) { + const int group_head = index / kTokenTile; + const int dim = dim_tile * kTokenTile + index % kTokenTile; + const int output_index = group_head * head_dim + dim; + shared_output[output_index] = + shared_output[output_index] * shared_old_scale[group_head] + + warp_pv[index]; + } + __syncwarp(); + } + __syncthreads(); + } + } + + for (int index = threadIdx.x; index < kGroupSize * head_dim; + index += blockDim.x) { + const int group_head = index / head_dim; + const int dim = index % head_dim; + const int query_head = kv_head * kGroupSize + group_head; + const int64_t scalar_offset = + (static_cast(token_id) * query_heads + query_head) * splits + + split_id; + partial_acc[scalar_offset * head_dim + dim] = shared_output[index]; + } + if (threadIdx.x < kGroupSize) { + const int query_head = kv_head * kGroupSize + threadIdx.x; + const int64_t scalar_offset = + (static_cast(token_id) * query_heads + query_head) * splits + + split_id; + partial_max[scalar_offset] = shared_row_max[threadIdx.x]; + partial_sum[scalar_offset] = shared_row_sum[threadIdx.x]; + } +} + +template +__global__ void InfLLMV2SparseAttentionSplitKVKernel( + const T* query, + const T* key_cache, + const T* value_cache, + const int* block_tables, + const int* seq_lens_decoder, + const int* batch_id_per_token, + const int* cu_seqlens_q, + const int* topk_indices, + float* partial_acc, + float* partial_max, + float* partial_sum, + int tokens, + int batch_size, + int physical_blocks, + int max_blocks_per_seq, + int query_heads, + int kv_heads, + int block_size, + int head_dim, + int capacity, + int splits, + int blocks_per_split) { + constexpr int kTokenTile = 8; + const int token_id = blockIdx.x; + const int kv_head = blockIdx.y; + const int split_id = blockIdx.z; + const int group_size = query_heads / kv_heads; + const int group_head = threadIdx.x / 32; + const int lane = threadIdx.x % 32; + if (token_id >= tokens || kv_head >= kv_heads || split_id >= splits || + group_head >= group_size) { + return; + } + const int query_head = kv_head * group_size + group_head; + const int values_per_lane = (head_dim + 31) / 32; + float query_values[8]; + float accumulator[8]; +#pragma unroll + for (int element = 0; element < 8; ++element) { + query_values[element] = 0.0f; + accumulator[element] = 0.0f; + } + const int64_t query_offset = + (static_cast(token_id) * query_heads + query_head) * head_dim; + for (int element = 0; element < values_per_lane; ++element) { + const int dim = lane + element * 32; + if (dim < head_dim) { + query_values[element] = InfLLMToFloat(query[query_offset + dim]); + } + } + + const int64_t partial_scalar_offset = + (static_cast(token_id) * query_heads + query_head) * splits + + split_id; + const int64_t partial_acc_offset = partial_scalar_offset * head_dim; + const int batch_id = batch_id_per_token[token_id]; + if (batch_id < 0 || batch_id >= batch_size) { + for (int element = 0; element < values_per_lane; ++element) { + const int dim = lane + element * 32; + if (dim < head_dim) { + partial_acc[partial_acc_offset + dim] = 0.0f; + } + } + if (lane == 0) { + partial_max[partial_scalar_offset] = -CUDART_INF_F; + partial_sum[partial_scalar_offset] = 0.0f; + } + return; + } + const int position = + InfLLMQueryPosition(token_id, batch_id, seq_lens_decoder, cu_seqlens_q); + const int* selected = + topk_indices + (token_id * kv_heads + kv_head) * capacity; + const float scale = rsqrtf(static_cast(head_dim)); + const int first_slot = split_id * blocks_per_split; + const int last_slot = min(capacity, first_slot + blocks_per_split); + + extern __shared__ char shared_bytes[]; + T* shared_key = reinterpret_cast(shared_bytes); + T* shared_value = shared_key + kTokenTile * head_dim; + float row_max = -CUDART_INF_F; + float row_sum = 0.0f; + for (int slot = first_slot; slot < last_slot; ++slot) { + const int logical_block = selected[slot]; + if (logical_block < 0) { + break; + } + const int physical_block = + block_tables[batch_id * max_blocks_per_seq + logical_block]; + if (physical_block < 0 || physical_block >= physical_blocks) { + continue; + } + for (int block_offset = 0; block_offset < block_size; + block_offset += kTokenTile) { + const int first_position = logical_block * block_size + block_offset; + if (first_position > position) { + break; + } + const int tile_tokens = + min(kTokenTile, + min(block_size - block_offset, position - first_position + 1)); + for (int index = threadIdx.x; index < tile_tokens * head_dim; + index += blockDim.x) { + const int tile_token = index / head_dim; + const int dim = index % head_dim; + const int64_t cache_offset = + ((static_cast(physical_block) * kv_heads + kv_head) * + block_size + + block_offset + tile_token) * + head_dim + + dim; + shared_key[index] = key_cache[cache_offset]; + shared_value[index] = value_cache[cache_offset]; + } + __syncthreads(); + + for (int tile_token = 0; tile_token < tile_tokens; ++tile_token) { + const int tile_offset = tile_token * head_dim; + float score = 0.0f; + for (int element = 0; element < values_per_lane; ++element) { + const int dim = lane + element * 32; + if (dim < head_dim) { + score += query_values[element] * + InfLLMToFloat(shared_key[tile_offset + dim]); + } + } + for (int offset = 16; offset > 0; offset >>= 1) { + score += __shfl_down_sync(0xffffffff, score, offset); + } + score = __shfl_sync(0xffffffff, score, 0) * scale; + float old_weight = 1.0f; + float new_weight = 1.0f; + if (score > row_max) { + old_weight = expf(row_max - score); + row_max = score; + } else { + new_weight = expf(score - row_max); + } + row_sum = row_sum * old_weight + new_weight; + for (int element = 0; element < values_per_lane; ++element) { + const int dim = lane + element * 32; + if (dim < head_dim) { + accumulator[element] = + accumulator[element] * old_weight + + new_weight * InfLLMToFloat(shared_value[tile_offset + dim]); + } + } + } + __syncthreads(); + } + } + for (int element = 0; element < values_per_lane; ++element) { + const int dim = lane + element * 32; + if (dim < head_dim) { + partial_acc[partial_acc_offset + dim] = accumulator[element]; + } + } + if (lane == 0) { + partial_max[partial_scalar_offset] = row_max; + partial_sum[partial_scalar_offset] = row_sum; + } +} + +template +__global__ void InfLLMV2SparseAttentionCombineKernel(const float* partial_acc, + const float* partial_max, + const float* partial_sum, + T* output, + int tokens, + int query_heads, + int head_dim, + int splits) { + const int token_id = blockIdx.x; + const int query_head = blockIdx.y; + const int thread = threadIdx.x; + if (token_id >= tokens || query_head >= query_heads) { + return; + } + const int64_t scalar_offset = + (static_cast(token_id) * query_heads + query_head) * splits; + extern __shared__ float reduction[]; + float local_max = -CUDART_INF_F; + for (int split = thread; split < splits; split += blockDim.x) { + local_max = fmaxf(local_max, partial_max[scalar_offset + split]); + } + reduction[thread] = local_max; + __syncthreads(); + for (int width = blockDim.x / 2; width > 0; width >>= 1) { + if (thread < width) { + reduction[thread] = fmaxf(reduction[thread], reduction[thread + width]); + } + __syncthreads(); + } + const float row_max = reduction[0]; + float local_sum = 0.0f; + if (row_max != -CUDART_INF_F) { + for (int split = thread; split < splits; split += blockDim.x) { + const float split_sum = partial_sum[scalar_offset + split]; + if (split_sum > 0.0f) { + local_sum += + split_sum * expf(partial_max[scalar_offset + split] - row_max); + } + } + } + reduction[thread] = local_sum; + __syncthreads(); + for (int width = blockDim.x / 2; width > 0; width >>= 1) { + if (thread < width) { + reduction[thread] += reduction[thread + width]; + } + __syncthreads(); + } + const float row_sum = reduction[0]; + const int64_t output_offset = + (static_cast(token_id) * query_heads + query_head) * head_dim; + for (int dim = thread; dim < head_dim; dim += blockDim.x) { + float value = 0.0f; + if (row_sum > 0.0f) { + for (int split = 0; split < splits; ++split) { + const float split_sum = partial_sum[scalar_offset + split]; + if (split_sum > 0.0f) { + const int64_t acc_offset = (scalar_offset + split) * head_dim + dim; + value += partial_acc[acc_offset] * + expf(partial_max[scalar_offset + split] - row_max); + } + } + value /= row_sum; + } + output[output_offset + dim] = InfLLMFromFloat(value); + } +} + +} // namespace fastdeploy diff --git a/custom_ops/setup_ops.py b/custom_ops/setup_ops.py index df349f8ed3a..142d4c5a8d0 100644 --- a/custom_ops/setup_ops.py +++ b/custom_ops/setup_ops.py @@ -349,6 +349,8 @@ def find_end_files(directory, end_str): "gpu_ops/reasoning_phase_token_constraint.cu", "gpu_ops/get_attn_mask_q.cu", "gpu_ops/mega_moe_pre_dispatch.cu", + "gpu_ops/infllmv2_attention/infllmv2.cu", + "gpu_ops/decoder_write_cache_with_rope.cu", ] sm_versions = get_sm_version(archs) # Some kernels in this file require SM75+ instructions. Exclude them when building SM70 (V100). @@ -562,7 +564,6 @@ def find_end_files(directory, end_str): "python utils/auto_gen_template_attention.py --config gpu_ops/decode_unified_attention/template_config.json --output gpu_ops/decode_unified_attention/template_instantiation/autogen" ) sources += ["gpu_ops/decode_unified_attention.cu"] - sources += ["gpu_ops/decoder_write_cache_with_rope.cu"] sources += find_end_files("gpu_ops/decode_unified_attention", ".cu") # Hopper optimized mla sources += find_end_files("gpu_ops/mla_attn", ".cu") diff --git a/docs/features/infllmv2_attention.md b/docs/features/infllmv2_attention.md new file mode 100644 index 00000000000..1f9ff2fd399 --- /dev/null +++ b/docs/features/infllmv2_attention.md @@ -0,0 +1,839 @@ +[简体中文](../zh/features/infllmv2_attention.md) + +# InfLLM-V2 Sparse Attention + +The `INFLLMV2_ATTN` backend provides MiniCPM4.1 with two-stage sparse attention during decode and long-context prefill. Stage 1 selects logical cache blocks independently for each request and KV head. Stage 2 maps the logical IDs through `block_tables` and reads the already-written paged K/V cache. + +Initial batch-1 prefill uses dense FlashAttention through `dense_len`, then shares one Stage 1 selection across each 128-token query tile. Each sparse tile is split into fully visible selected history blocks and its one or two current blocks. The history partition uses non-causal FlashAttention, the current partition uses causal FlashAttention, and their outputs are combined with the two exact log-sum-exp values. This preserves causal alignment even though Paddle FlashAttention V2 aligns a shorter causal query at the upper left. Short, mixed, and shared chunked-prefill requests retain the dense path while still building the two-scale semantic summaries needed by decode. + +## Implementation overview + +The decode data flow is: + +```text +raw fused QKV + -> decoder_write_cache_with_rope + -> post-RoPE / post-QK-norm Q + -> current K/V committed to the paged cache + -> infllmv2_update_compressed_k + -> infllmv2_select_blocks + -> infllmv2_attention_forward + -> block_tables[request, logical block] -> physical page + -> causal sparse softmax +``` + +Both sparse stages consume the post-RoPE query returned by the writer, and Stage 2 observes the current decode token in the K/V cache. Raw fused QKV is never passed directly to sparse attention. + +The sparse prefill data flow is: + +```text +raw fused QKV + -> gqa_rope_write_cache + -> post-RoPE / post-QK-norm Q/K/V and paged K/V cache + -> infllmv2_update_compressed_k + -> dense FlashAttention for [0, dense_len) + -> one infllmv2_select_blocks call per 128-token query tile + -> gather selected paged K/V + -> non-causal FlashAttention over selected history pages + -> causal FlashAttention over the current one or two pages + -> exact LSE-weighted output merge +``` + +The default 4,096-token processing chunk contains 32 query tiles and bounds temporary gathered K/V storage. A final partial tile is supported, including prompts whose length is not block aligned. + +For every query token, the request and causal position come from runtime metadata: + +```text +request = batch_id_per_token[token] +local_offset = token - cu_seqlens_q[request] +position = seq_lens_decoder[request] + local_offset +``` + +This keeps requests, KV heads, and page tables isolated under continuous batching. Padding tokens with `batch_id_per_token == -1` do not access the cache. + +### Two-scale Stage 1 + +Stage 1 maintains two K-summary scales: + +| Scale | Default window | Default stride | Purpose | +| --- | ---: | ---: | --- | +| Fine | 32 | 16 | Estimate detailed relevance for logical KV blocks | +| Coarse | 128 | 64 | Approximate the log-sum-exp normalizer for each query head | + +Each summary is the mean of K vectors in one complete window. A completed window is owned by the physical page containing its final token, which preserves windows crossing page boundaries. With the default `block_size=64`, each physical block and KV head has four fine slots and one coarse slot. + +Let `G(g)` be the GQA query-head group associated with KV head `g`, and let `d` be the head dimension. Scores are computed as follows: + +```text +LSE[h] = logsumexp_j(dot(Q[h], Kbar_coarse[j]) / sqrt(d)) + +semantic_score[i, g] = sum(h in G(g)) exp( + dot(Q[h], Kbar_fine[i]) / sqrt(d) - LSE[h] +) + +block_score[b, g] = max(semantic_score[i, g] + for fine window i overlapping block b) +``` + +For the MiniCPM4.1 production shape (`QH=32`, `KVH=2`, `D=128`, +`block_size=64`), the coarse normalizer uses a GQA-tiled split kernel adapted +from the distributed-softmax structure in the OpenBMB CUDA implementation. +One CTA owns `(query token, KV head, coarse split)`. Its 16 half-warps keep the +16 Q heads in registers, cooperatively load one aligned `float4` K tile for up +to 16 coarse windows, and reuse that tile across the whole GQA group. The +per-split online `(max, sum)` values stay distributed across lanes; a second +one-warp-per-query-head kernel performs the rescaled split merge. The split +partials are caller-provided persistent workspaces, so this path has neither a +thread-0 window loop nor a per-call allocation. + +The production fine block-score path is tiled across eight adjacent candidate +blocks. One CTA owns `(query token, KV head, eight-block tile)`; its 16 +half-warps keep the GQA query heads in registers while up to 33 overlapping +fine K rows are loaded once with aligned `float4` transactions and shared by +the whole group. The same head contributions are then pooled for all eight +candidate blocks. Every score slot, including a partially filled table tail, +is rewritten because the backend reuses this workspace across decode steps. + +The first `init_blocks`, the preceding blocks covered by `window_size`, and the current block receive positive-infinity scores and participate in the same top-k; remaining candidates are ranked by `block_score`. With `local_blocks = window_size / block_size`, the sparse-region selection budget is `topk + local_blocks`, so `topk` is not an extra budget beyond the forced blocks. The default selects 96 blocks: 34 forced blocks (one initial, 32 preceding local, and the current block) and at most 62 ordinary dynamic blocks. The output buffer also has to accommodate the select-all behavior below the dense threshold, so its capacity is: + +```text +selected_capacity = max( + topk + window_size / block_size, + ceil(dense_len / block_size), +) +``` + +Sparse-region selection uses an exact 64-bit composite radix key: score +descending, then logical block ID ascending. A 256-thread CUB block radix sort +dispatches 1, 2, 4, or 8 candidates per thread and therefore covers up to +2,048 logical blocks (128K tokens at `block_size=64`) in one CTA. This replaces +the previous all-pairs rank computation while preserving deterministic ties. +Requests below `dense_len` bypass sorting and emit all visible blocks directly. +Model configurations exceeding 2,048 candidates fail explicitly. + +Stage 1 returns GPU tensors with the following contracts: + +| Tensor | Shape | Meaning | +| --- | --- | --- | +| `topk_indices` | `[tokens, kv_heads, selected_capacity]` | Logical block IDs per request/query token/KV head; unused slots are `-1` | +| `block_scores` | `[tokens, kv_heads, max_blocks_per_seq]` | Dynamic block scores | +| `selected_counts` | `[tokens, kv_heads]` | Number of selected blocks | + +The valid prefix of `topk_indices` must contain strictly increasing, unique logical block IDs within the request's visible range; every remaining slot must be `-1`. Stage 2 uses the first `-1` to terminate the prefix. This is a custom-op precondition guaranteed by Stage 1 in normal serving. FastDeploy's scheduler must likewise provide self-consistent sequence lengths, `cu_seqlens_q`, token-to-request mappings, and `block_tables` page ownership; active requests must not alias writable physical pages. Q/K/V inputs must contain finite values. + +MiniCPM4.1 uses GQA, so metadata is produced per KV head rather than duplicated for every query head in its GQA group. Stage 2 shares that KV-head block list across the associated query heads. + +### Paged Stage 2 + +Stage 2 does not assume that logical blocks are contiguous in memory. Every selected block is resolved using: + +```text +physical_block = block_tables[request, logical_block] +``` + +Scaled dot-product attention is then evaluated only for tokens in that physical page that do not exceed the request's current logical `position`. Values in `topk_indices` must therefore be logical block IDs, not physical page IDs. + +For the MiniCPM4.1 production shape (`32` query heads, `2` KV heads, `head_dim=128`, `block_size=64`), Stage 2 uses a GQA-tiled FlashDecoding kernel. One 128-thread CTA owns `(query token, KV head, KV split)` and covers all 16 query heads sharing that KV head. It retains a 16-token K/V shared-memory tile: warp 0 evaluates the 16-by-16 QK tensor-core tile, threads 0 through 15 update the per-head online `(max, sum)` state, and all four warps evaluate the probability-times-V tensor-core tiles and update the shared FP32 accumulators. A full 16-token K tile and V tile are copied from the paged cache to shared memory with aligned 16-byte `uint4` transactions; a causal boundary tile uses the explicit scalar masked path. A separate kernel merges the split `(accumulator, max, sum)` partials. + +Each KV split contains two selected pages. This yields enough CTAs for batch-1 decode without restoring the old 16-fold K/V global-memory reads. The local CUDA repository instead uses a FlashAttention/CUTE 64-token K/V tile. FastDeploy keeps the smaller tile because a directly ported 64-token paged tile needs 47.3 KiB shared memory and regressed the 128K/batch-4 Stage 2 measurement by 17.1%; the 16-token paged tile uses about 22 KiB while preserving GQA reuse, split-KV, and online softmax. Non-production shapes use the generic correctness path; they do not represent the MiniCPM4.1 performance path. + +Each layer backend owns persistent tensors for Stage 1 metadata, final attention output, and the FP32 split `(accumulator, max, sum)` values. All are passed as in-place custom-op inputs and reused across decode steps with the same shape. The CUDA launcher performs no `paddle::empty`, `paddle::full`, or `paddle::zeros` allocation on the Stage 2 hot path. + +## Configuration and activation + +Select the backend explicitly: + +```shell +export FD_ATTENTION_BACKEND=INFLLMV2_ATTN +``` + +Sparse parameters are read first from `sparse_config` in the model's `config.json`, then from same-named top-level model fields, and finally from the MiniCPM4.1 defaults. To override them, add: + +```json +{ + "sparse_config": { + "kernel_size": 32, + "kernel_stride": 16, + "topk": 64, + "dense_len": 8192, + "init_blocks": 1, + "window_size": 2048, + "sparse_prefill": true, + "prefill_query_chunk_size": 4096 + } +} +``` + +| Parameter | Default | Constraint | +| --- | ---: | --- | +| `block_size` | 64 | The `--block-size` server option; divisible by both `kernel_stride` and `4 * kernel_stride` | +| `kernel_size` | 32 | Positive integer | +| `kernel_stride` | 16 | Positive integer | +| `topk` | 64 | Positive integer | +| `dense_len` | 8192 | At least `4 * kernel_size` | +| `init_blocks` | 1 | Non-negative integer | +| `window_size` | 2048 | A non-negative multiple of `block_size` | +| `sparse_prefill` | `true` | Enables the model-specific initial-prefill sparse path | +| `prefill_query_chunk_size` | 4096 | Positive multiple of 128 | + +`init_blocks` must be smaller than `topk` so the initial blocks and the inclusive local set fit in the output capacity. + +The following example starts a single-GPU server for a 32K context: + +```shell +export CUDA_VISIBLE_DEVICES=0 +export FD_ATTENTION_BACKEND=INFLLMV2_ATTN +export MODEL_PATH=/path/to/MiniCPM4.1-8B + +.venv/bin/python -m fastdeploy.entrypoints.openai.api_server \ + --model "${MODEL_PATH}" \ + --served-model-name MiniCPM4.1-8B \ + --port 8180 --metrics-port 8181 \ + --engine-worker-queue-port 8182 --cache-queue-port 8183 \ + --tensor-parallel-size 1 \ + --max-model-len 65536 \ + --block-size 64 \ + --max-num-seqs 1 \ + --max-num-batched-tokens 32768 \ + --gpu-memory-utilization 0.5 \ + --no-enable-prefix-caching \ + --graph-optimization-config '{"use_cudagraph": false}' +``` + +## Limitations and memory cost + +- The custom operators support FP32, FP16, and BF16 on NVIDIA CUDA devices with compute capability 8.0 or newer (SM80+). +- The paged K/V cache must be unquantized and rank 4. Quantized KV cache fails explicitly. +- The number of query heads must be divisible by the number of KV heads, and `head_dim` must be in `[1, 256]`. +- Tensor parallelism must not replicate KV heads: `tensor_parallel_size` must be no larger than every layer's global KV-head count. Replica-spanning Stage 1 score reduction is not implemented yet, so unsupported configurations fail explicitly. +- Sparse decode requires semantic K summaries initialized by prefill on the same paged K cache. P/D disaggregation does not currently transfer these summaries and must not enter sparse decode without them. +- Resetting or replacing the paged K cache must also reset its semantic summaries. +- Semantic summaries do not yet participate in CUDA Graph cache replacement, so CUDA Graph must currently be disabled. +- Speculative decoding/MTP is rejected until every proposed token participates in Stage 1 metadata and semantic-cache updates. +- Sparse prefill currently activates only for an initial causal batch-1 request with the MiniCPM4.1 production shape (`QH=32`, `KVH=2`, `D=128`, `block_size=64`) and FP16/BF16 cache. Other layouts fall back to dense prefill. +- The 32K acceptance run uses one scheduler-level 32K prefill request (internally processed in 4,096-token sparse batches) and therefore does not exercise scheduler-level multi-chunk prefill. With `--max-num-batched-tokens 8192 --enable-chunked-prefill`, both the existing FlashAttention service and InfLLM-V2 service abort in the third 8K prefill chunk on an uninitialized Paddle `optional::get()`. This isolates the failure to the shared chunked-prefill path rather than sparse Stage 1/Stage 2. This implementation does not change that shared path; both comparison variants use `--max-num-batched-tokens 32768`. + +The theoretical additional semantic-cache storage per layer and physical block is: + +```text +dtype_bytes * kv_heads * head_dim * ( + block_size / kernel_stride + + block_size / (4 * kernel_stride) +) +``` + +For MiniCPM4.1 BF16 with `kv_heads=2`, `head_dim=128`, and the default windows, this is 2,560 bytes/block/layer. This is a shape-derived static cost, not the measured INF-3 peak memory result. + +## Build and regression + +The CUDA custom operators must be rebuilt before operator tests are run. This example uses the project's `.venv`; adjust `CUDA_HOME` and the compute-capability list for the target machine: + +```shell +export CUDA_HOME=/path/to/cuda +export PATH="${CUDA_HOME}/bin:${PATH}" +export LD_LIBRARY_PATH="${CUDA_HOME}/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" +export PYTHONPATH="${PWD}" + +MAX_JOBS=32 FD_BUILD_RESUME=1 \ + bash build.sh 0 "${PWD}/.venv/bin/python" false "[86]" +``` + +Run the focused regressions with: + +```shell +CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. .venv/bin/python -m pytest -q \ + tests/model_executor/test_infllmv2_attention_backend.py \ + tests/operators/test_infllmv2_attention_forward.py +``` + +The backend tests cover writer/Stage 1/Stage 2 ordering, post-RoPE query use, and semantic-cache lifetime. GPU operator tests cover cross-page semantic windows, per-request/per-KV-head selection, paged mapping, the short-context closed loop, and invalid metadata. This feature does not require documentation tests and must not modify `tests/model_executor/test_minicpm41.py`. + +## Dense/sparse long-context comparison + +`benchmarks/benchmark_infllmv2.py` has seven subcommands: + +- `selector` inspects a raw selector trace and writes an unbound hit-rate summary. This summary is for inspection only and cannot be attached to a timed sparse run. +- `sparse-diagnostic` runs the sparse workload once against a trace-enabled server, automatically forwards `--no-warmup` to `benchmark_serving.py`, and binds the newly generated trace to the exact prompt-token workload in a selector diagnostic JSON. +- `run` delegates throughput, TTFT, and TPOT collection to `benchmark_serving.py` while sampling memory on one GPU. A timed sparse run requires the bound JSON through `--selector-diagnostic`. +- `report` produces a comparison JSON and Markdown table only when dense and sparse workload fingerprints match exactly and both runs report the same physical GPU UUID. +- `operators` uses CUDA events to time compressed-K update, Stage 1, Stage 2, and the complete sparse decode operator chain. Its default matrix is 32K/concurrency 1 and 128K/concurrency 4. +- `prefill` compares the complete dense and sparse Paddle prefill paths with CUDA events. +- `cuda-impl` runs the equivalent decode workload against a local `infllm_v2` PyTorch extension checkout. + +Run the calibrated operator matrix independently from service timing: + +```shell +CUDA_VISIBLE_DEVICES=7 .venv/bin/python benchmarks/benchmark_infllmv2.py operators \ + --gpu-index 7 \ + --scenario 32768:1 \ + --scenario 131072:4 \ + --output runs/bench/infllmv2/operators.json +``` + +The JSON records mean/median/P10/P90/P99 CUDA time, host wall time, source hashes, the selected-token fraction, and persistent workspace bytes for each scenario. It also verifies that all four Stage 2 outputs alias the caller-provided buffers. This synthetic benchmark isolates kernel costs; it is not a substitute for TTFT/TPOT service comparison. + +Sparse prefill has a dedicated complete-chain benchmark: + +```shell +CUDA_VISIBLE_DEVICES=7 PYTHONPATH=. .venv/bin/python \ + benchmarks/benchmark_infllmv2.py prefill \ + --device gpu:0 \ + --context-length 16384 \ + --context-length 32768 \ + --output runs/bench/infllmv2/sparse_prefill.json +``` + +The sparse timing includes compressed-K update, the dense prefix, Stage 1, paged-cache gather, both FlashAttention partitions, and the LSE merge. On an NVIDIA RTX A6000 with BF16 and batch 1, the 2026-08-18 run measured 90.300 ms dense versus 46.005 ms sparse at 32K, or **1.963x**. At 16K it measured 21.256 ms versus 18.654 ms, or 1.139x, showing that the optimization is intended for long contexts. The traceable result is `runs/bench/infllmv2/sparse_prefill_20260818.json`. + +### 32K/batch-1 theoretical ceiling + +With defaults, sparse Stage 2 selects 96 pages of 64 tokens. At 32K this is 6,144 of 32,768 visible tokens, or 18.75%. Even under an ideal bandwidth-only model, Stage 2's token-reduction ceiling is therefore `32768 / 6144 = 5.333x`. End-to-end decode cannot exceed that number and is bounded more tightly by Amdahl's law: + +```text +speedup <= 1 / ((1 - f_stage2) + f_stage2 / 5.333) +``` + +where `f_stage2` is dense decode's Stage 2-attention time fraction. Stage 1, compressed-K update, split merge, launch overhead, MLP/MoE, and batch-1 GPU underfill are outside the idealized reduction and lower the observed gain. At 128K the selected-token ratio is lower, but Stage 1 scans more semantic windows, so the token ratio alone must not be reported as measured speedup. + +### Selector hit rate + +Hit rate is the micro-averaged dense-reference block recall across all request/query-token/KV-head samples: + +```text +sum(|selected_blocks ∩ reference_blocks|) / sum(|reference_blocks|) +``` + +`reference_blocks` should come from exact dense attention using the same post-RoPE query and paged K cache, aggregated over each GQA group to KV-head granularity. The raw format is: + +```json +{ + "samples": [ + { + "selected_blocks": [0, 12, 37], + "reference_blocks": [0, 9, 37] + } + ] +} +``` + +Block IDs must be non-negative and unique within an array, and `reference_blocks` must not be empty. These numbers only demonstrate the JSON format; they are not measured model results. + +The backend can produce these samples directly in a separate diagnostic run. Set the trace variables before starting a dedicated `INFLLMV2_ATTN` diagnostic server: + +```shell +export FD_INFLLMV2_SELECTOR_TRACE_PATH="${PWD}/runs/bench/infllmv2/selector_samples.json" +export FD_INFLLMV2_SELECTOR_TRACE_RANK=0 +export FD_INFLLMV2_SELECTOR_TRACE_LAYER=0 +export FD_INFLLMV2_SELECTOR_TRACE_MAX_SAMPLES=16 +``` + +CUDA Graph must be disabled while tracing. After Stage 1, the backend computes an FP32 dense softmax from the same post-RoPE query, `block_tables`, and already-written paged K cache. It sums logical-block attention mass over the query heads sharing each KV head and uses the highest-mass blocks, with the same cardinality as `selected_count`, as the reference. Tracing introduces GPU synchronization, device-to-host copies, and dense-attention work, so diagnostic latency is never treated as a performance result. The trace path must not exist before the diagnostic, and the server refuses to overwrite it. + +```shell +mkdir -p runs/bench/infllmv2 + +.venv/bin/python benchmarks/benchmark_infllmv2.py sparse-diagnostic \ + --base-url http://127.0.0.1:8180 \ + --model MiniCPM4.1-8B \ + --tokenizer "${MODEL_PATH}" \ + --input-len 32768 \ + --output-len 128 \ + --num-prompts 4 \ + --max-concurrency 1 \ + --trace-path runs/bench/infllmv2/selector_samples.json \ + --output runs/bench/infllmv2/selector_diagnostic.json +``` + +`sparse-diagnostic` automatically disables the normal warm-up request by passing `--no-warmup`; there is no separate warm-up result to report. It records the exact generated prompt-token digest and workload fingerprint together with the trace provenance. The later timed sparse `run` rejects a diagnostic whose workload differs. + +The raw trace can also be summarized for standalone inspection, but this unbound output is not accepted by `run`: + +```shell +.venv/bin/python benchmarks/benchmark_infllmv2.py selector \ + --samples runs/bench/infllmv2/selector_samples.json \ + --output runs/bench/infllmv2/selector_summary.json +``` + +### Fixed-workload runs + +After the diagnostic finishes, stop the trace-enabled sparse server completely. Keep the raw trace file because the bound diagnostic verifies it by path and SHA256, but disable tracing and restart the server before any timing measurement: + +```shell +unset FD_INFLLMV2_SELECTOR_TRACE_PATH +``` + +The diagnostic server and timed server must be separate processes. Confirm that the timed `INFLLMV2_ATTN` server starts without the trace-path variable; trace instrumentation must remain disabled for the complete timed run. + +For the timed comparison, first start a dense server with `FD_ATTENTION_BACKEND=FLASH_ATTN`, then run: + +```shell +.venv/bin/python benchmarks/benchmark_infllmv2.py run \ + --variant dense \ + --base-url http://127.0.0.1:8180 \ + --model MiniCPM4.1-8B \ + --tokenizer "${MODEL_PATH}" \ + --gpu-index 0 \ + --input-len 32768 \ + --output-len 128 \ + --num-prompts 4 \ + --max-concurrency 1 \ + --output runs/bench/infllmv2/dense.json +``` + +Stop the dense server. Start the untraced `INFLLMV2_ATTN` server on the same idle GPU with otherwise identical server arguments, then run: + +```shell +.venv/bin/python benchmarks/benchmark_infllmv2.py run \ + --variant sparse \ + --base-url http://127.0.0.1:8180 \ + --model MiniCPM4.1-8B \ + --tokenizer "${MODEL_PATH}" \ + --gpu-index 0 \ + --input-len 32768 \ + --output-len 128 \ + --num-prompts 4 \ + --max-concurrency 1 \ + --selector-diagnostic runs/bench/infllmv2/selector_diagnostic.json \ + --output runs/bench/infllmv2/sparse.json + +.venv/bin/python benchmarks/benchmark_infllmv2.py report \ + --dense-result runs/bench/infllmv2/dense.json \ + --sparse-result runs/bench/infllmv2/sparse.json \ + --output runs/bench/infllmv2/report.json +``` + +The workload arguments supplied to `sparse-diagnostic`, dense `run`, and sparse `run` must be identical, including model, tokenizer, endpoint, seed, prompt count, input/output lengths, request rate, and concurrency. The script verifies the generated prompt tokens and rejects mismatched fingerprints. + +`--gpu-index` is the physical index consumed by `nvidia-smi --id`. If the server uses `CUDA_VISIBLE_DEVICES=7`, pass `--gpu-index 7`, not the process-local index 0. Keep other workloads off that GPU and restart the server completely between variants. + +In each `run` result, `gpu_memory.gpu_uuid` identifies the physical device, `baseline_mib` is the already-started server's memory immediately before requests, `peak_mib` is total used GPU memory at peak, and `peak_delta_mib` is their difference. `report` rejects dense and sparse inputs with different GPU UUIDs and preserves the UUID plus baseline, peak, and peak-delta values in the report. An acceptance report must retain all of them. + +### Retained P0/P2 checkpoint results + +The final CUDA-event matrix is stored in `runs/bench/infllmv2/tmp/p3_operator_matrix.json`. It uses the MiniCPM4.1 production shape (`QH=32`, `KVH=2`, `D=128`, BF16), 50 warm-up iterations, and 200 measured iterations on an NVIDIA RTX A6000. Values below are median device times. + +| Scenario | Compressed-K update (us) | Stage 1 (us) | Stage 2 (us) | Complete sparse chain (us) | Stage 2 persistent bytes | +| --- | ---: | ---: | ---: | ---: | ---: | +| 32K / concurrency 1 | 49.152 | 191.488 | 111.616 | 380.928 | 1,073,152 | +| 128K / concurrency 4 | 49.152 | 1,662.976 | 215.040 | 1,939.968 | 4,292,608 | + +The pre-P0/P2 Stage 2 baseline is retained in `runs/bench/infllmv2/tmp/p0_p2_operator_baseline.json`. Comparing the same shapes on the same GPU gives: + +| Scenario | Pre-P0 Stage 2 (us) | GQA-tiled Stage 2 (us) | Speedup | Time reduction | +| --- | ---: | ---: | ---: | ---: | +| 32K / concurrency 1 | 135.168 | 111.616 | 1.211x | 17.42% | +| 128K / concurrency 4 | 367.616 | 215.040 | 1.710x | 41.50% | + +The dedicated P0 checkpoint itself measured 112.144 us and 215.040 us for Stage 2, or 1.205x and 1.710x over the retained baseline. The 32K difference from the final 111.616 us value is run-to-run variation. + +The matrix also verifies that the output and all three FP32 partial tensors alias caller-provided storage. These byte counts are per layer-backend instance: Stage 2 retains 1,073,152 bytes at concurrency 1 and 4,292,608 bytes at concurrency 4. Each instance allocates on shape change rather than on every layer invocation; for this 32-layer checkpoint, the corresponding totals are 32.75 MiB and 131.00 MiB. The P2 checkpoint measured 112.640 us for 32K/1 Stage 2 versus 112.144 us at the P0 checkpoint (+0.44%), while the full chain changed from 381.952 us to 380.928 us (-0.27%); 128K/4 Stage 2 remained 215.040 us. This is within run-to-run noise and shows no hot-path regression from in-place persistence. The P0 and P2 checkpoints are preserved separately as `p0_gqa_tiled_operator.json` and `p2_persistent_workspace_operator.json` in the same temporary-results directory. + +The checkpoint used here declares both `max_position_embeddings` and `rope_scaling.original_max_position_embeddings` as 65,536. A 128K service request would violate that model contract, so the required 128K/concurrency-greater-than-one control is intentionally a synthetic operator comparison and is not reported as an end-to-end model result. + +### Stage 1 selector optimizations + +The subsequent Stage 1 GQA-tiled coarse-LSE and exact radix-TopK measurements +are stored separately in +`runs/bench/infllmv2/tmp/stage1_p1_gqa_tiled_result.json` and +`runs/bench/infllmv2/tmp/stage1_p2_radix_topk_result.json`. The combined CUDA +event result is `stage1_p1_p2_optimized.json`; Nsight Systems traces and SQLite +exports in the same directory provide the per-kernel attribution. These runs +use the same RTX A6000 and BF16 production shape, with 20 warm-ups and 100 +measured iterations for CUDA-event medians. + +| Scenario | Baseline Stage 1 (us) | Optimized Stage 1 (us) | Stage 1 speedup | Baseline chain (us) | Optimized chain (us) | Chain speedup | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| 32K / concurrency 1 | 191.488 | 82.944 | 2.309x | 381.952 | 246.784 | 1.548x | +| 128K / concurrency 4 | 1,659.904 | 820.224 | 2.024x | 1,935.360 | 1,052.672 | 1.839x | + +| Scenario | Baseline coarse LSE (us) | GQA-tiled coarse LSE (us) | Speedup | Baseline TopK (us) | Radix TopK (us) | Speedup | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| 32K / concurrency 1 | 111.155 | 7.635 | 14.558x | 34.643 | 13.978 | 2.478x | +| 128K / concurrency 4 | 479.877 | 22.688 | 21.151x | 509.055 | 22.579 | 22.545x | + +At that checkpoint, block-score calculation was the dominant Stage 1 kernel at +128K/concurrency 4 (773.271 us, about 94% of Stage 1) and was therefore the +next selector bottleneck. Its isolated extension passed 10 SM86 GPU +correctness cases, including a production-shape BF16 semantic reference and +2,048-candidate exact-tie/padding cases; the backend regression suite passed 39 +cases. After that checkpoint's package rebuild, the packaged SM86 operator +suite passed all 15 cases. The block-score bottleneck is addressed in the next +section. + +### GQA-tiled block-score and final Stage 2 refinements + +The follow-up implementation was informed by the split-KV and distributed +softmax structure in the official OpenBMB repository, cloned under +`runs/bench/infllmv2/tmp/infllmv2_cuda_impl` at commit +`93cf2ec28e5a7acebe3f0bb7329b6c73a1be91f6`. FastDeploy retains its own paged +cache, persistent-workspace, and in-place custom-op contracts. + +The final isolated CUDA-event matrix is +`runs/bench/infllmv2/tmp/stage1_p3_gqa_blockscore/operator_optimized_halfwarp.json`. +It uses the same RTX A6000, BF16 production shape, 20 warm-ups, and 100 measured +iterations as the preceding selector comparison. + +| Scenario | Compressed-K update (us) | Stage 1 (us) | Stage 2 (us) | Complete sparse chain (us) | +| --- | ---: | ---: | ---: | ---: | +| 32K / concurrency 1 | 49.152 | 48.128 | 59.392 | 149.504 | +| 128K / concurrency 4 | 49.152 | 131.072 | 124.928 | 317.440 | + +Relative to the retained P0/P2 matrix, the complete chain falls from 380.928 +to 149.504 us at 32K/1 (2.548x, 60.75% less time) and from 1,939.968 to +317.440 us at 128K/4 (6.111x, 83.64% less time). Relative to the immediately +preceding coarse-LSE/radix-TopK checkpoint, Stage 1 is 1.723x faster at 32K/1 +and 6.258x faster at 128K/4; Stage 2 is 1.879x and 1.787x faster. + +Nsight Systems reports, SQLite exports, and generated kernel-summary CSV files +are in the same temporary directory. The per-kernel averages below compare the +preceding checkpoint with the final implementation. + +| Kernel | 32K before/final (us) | Speedup | 128K/4 before/final (us) | Speedup | +| --- | ---: | ---: | ---: | ---: | +| Fine block-score | 58.585 / 13.370 | 4.382x | 773.271 / 82.341 | 9.391x | +| Stage 2 partial | 70.386 / 40.153 | 1.753x | 152.721 / 95.404 | 1.601x | +| Stage 2 merge | 37.548 / 15.059 | 2.493x | 72.025 / 25.997 | 2.770x | + +The current-source isolated extension passed 14/14 GPU cases, including +FP32/FP16/BF16 paged Stage 2, production BF16 GQA Stage 1/2, exact 2,048-way +ties, reusable-buffer tails, and invalid contracts. The production package and +matched service result are reported after the full rebuild below. + +For the supported 32K service workload, the retained P0/P2 checkpoint report +`runs/bench/infllmv2/tmp/p3_report_32k_current.json` compares dense and sparse +servers launched from that same packaged build with identical service +parameters. It predates the Stage 1 selector optimizations above. Both inputs +have workload ID +`b0fd4419f876adfa40e4d92a6e46d56c95c0f5f37e0539abd1a32c8208f6c078`, the +same prompt-token hash, and the same physical GPU UUID; the report command +validated these fields. + +| Metric | Dense | P0/P2 checkpoint sparse | Sparse / Dense | +| --- | ---: | ---: | ---: | +| Selector block hit rate | — | 64.4531% (990 / 1536) | — | +| Request throughput (req/s) | 0.092799 | 0.084076 | 0.906x | +| Output throughput (token/s) | 11.878301 | 10.761664 | 0.906x | +| Total-token throughput (token/s) | 3052.723423 | 2765.747673 | 0.906x | +| Mean / median / P99 TTFT (ms) | 7368.456 / 7362.350 / 7399.758 | 7404.749 / 7389.901 / 7467.811 | 1.005x / 1.004x / 1.009x | +| Mean / median / P99 TPOT (ms) | 26.742 / 26.744 / 26.794 | 35.252 / 35.408 / 35.514 | 1.318x / 1.324x / 1.325x | +| Baseline / peak / peak-delta GPU memory (MiB) | 19493 / 24421 / 4928 | 19417 / 24345 / 4928 | 0.996x / 0.997x / 1.000x | + +The service workload contains four deterministic random-token requests with input 32768, output 128, concurrency 1, request rate `inf`, and seed 2026. It uses TP1, BF16, `block_size=64`, default sparse configuration `32/16/64/8192/1/2048`, no prefix cache, no CUDA Graph, and one 32K prefill chunk. The environment is an NVIDIA RTX A6000 48 GiB (UUID `GPU-56bb1092-c218-9cd3-ad85-9a2735522d28`), CUDA toolkit 12.8, Paddle 3.3.1 built with CUDA 12.6, and Python 3.12.13. + +That retained service result is not an end-to-end measurement of the new Stage +1 kernels: its sparse output throughput is 0.906x dense and mean TPOT is 1.318x +dense. The current-source operator chain is 35.39% shorter at 32K/1 and 45.61% +shorter at 128K/4. At that checkpoint, a packaged server rebuild and matched +service rerun were still required; the completed rerun is reported below. The older +`runs/bench/infllmv2/report.json` remains a still earlier historical result +(0.311x dense output throughput and 7.945x dense mean TPOT) and must not be +presented as the current kernel result. + +### Historical post-Stage1 packaged service checkpoint + +The full `fastdeploy_ops` package was rebuilt from the then-current sources +before this checkpoint. The installed library had SHA-256 +`cb37a51292ed2f92139a199f748a1e8fb014d6c7774284cfdc6bfe535ed2d5e4`, and +the packaged GPU operator regression passed 15/15 cases. Build, correctness, +server, and benchmark logs are retained under +`runs/bench/infllmv2/tmp/retest_stage1_20260813/`. + +The validated comparison is +`runs/bench/infllmv2/tmp/retest_stage1_20260813/report_32k.json`. Dense and +sparse use the same workload ID and prompt-token hash shown above, the same +physical GPU UUID, and completely restarted servers. The timed sparse server +did not have selector tracing enabled. + +| Metric | Dense | Post-Stage1 checkpoint sparse | Sparse / Dense | +| --- | ---: | ---: | ---: | +| Selector block hit rate | — | 59.2448% (910 / 1536) | — | +| Request throughput (req/s) | 0.092539 | 0.087938 | 0.950x | +| Output throughput (token/s) | 11.845043 | 11.256072 | 0.950x | +| Total-token throughput (token/s) | 3044.176118 | 2892.810553 | 0.950x | +| Mean / median / P99 TTFT (ms) | 7403.585 / 7401.545 / 7445.184 | 7414.495 / 7413.212 / 7430.486 | 1.001x / 1.002x / 0.998x | +| Mean / median / P99 TPOT (ms) | 26.710 / 26.717 / 26.747 | 31.071 / 31.058 / 31.133 | 1.163x / 1.162x / 1.164x | +| Baseline / peak / peak-delta GPU memory (MiB) | 19493 / 24421 / 4928 | 19417 / 24345 / 4928 | 0.996x / 0.997x / 1.000x | + +Relative to the retained P0/P2 sparse checkpoint, mean sparse TPOT fell from +35.252 ms to 31.071 ms (11.86% reduction, 1.135x speedup) and sparse output +throughput rose from 10.761664 to 11.256072 token/s (+4.59%). The repeated +dense mean TPOT changed by only -0.12%, so the sparse gain is not explained by +a faster dense baseline. Selector diagnostics follow generated decode +trajectories; the diagnostic runs share prompt fingerprints but not generated +token identities, so the current 59.24% and historical 64.45% hit rates should +not be treated as a controlled selector-quality delta. + +At this historical checkpoint, sparse improved but did not cross dense at +32K/batch 1: its mean TPOT +is still 4.361 ms higher (1.163x dense), and output throughput is 0.950x dense. +The 5.333x token-reduction ceiling applies only to Stage 2 K/V work. Dense +FlashAttention is already fused and efficient, while sparse decode additionally +pays compressed-K update, block-score/coarse-LSE/TopK selection, split merge, +kernel-launch overhead across 32 layers, and the same MLP/MoE cost. Batch 1 also +underfills the GPU. These fixed costs exceed the saved Stage 2 K/V traffic in +this checkpoint, so sparsity alone does not imply an end-to-end speedup. + +### Final GQA-tiled packaged validation and service rerun + +After the block-score and final Stage 2 changes, the complete SM86 production +package was rebuilt again. The repository-packaged +`fastdeploy/model_executor/ops/gpu/fastdeploy_ops/fastdeploy_ops_pd_.so` has +SHA-256 +`368174499c47eb991595616988f1c2a65bd6384f8cf8353d366511784883bb94`. +The backend regression passed 39/39 cases and the packaged GPU operator suite +passed 15/15 cases. The successful build and test logs are under +`runs/bench/infllmv2/tmp/stage1_p3_gqa_blockscore/full_build/`. + +The packaged CUDA-event rerun is +`runs/bench/infllmv2/tmp/final_gqa_decode_20260813/operators_packaged.json`. +It records the final source hashes, verifies all four Stage 2 output aliases, +and gives these device medians: + +| Scenario | Compressed-K update (us) | Stage 1 (us) | Stage 2 (us) | Complete sparse chain (us) | +| --- | ---: | ---: | ---: | ---: | +| 32K / concurrency 1 | 49.152 | 37.888 | 59.392 | 149.504 | +| 128K / concurrency 4 | 49.152 | 131.072 | 125.952 | 316.416 | + +Against the retained P0/P2 chain this is a 2.548x speedup at 32K/1 and a +6.131x speedup at 128K/4. Against the pre-P0 Stage 2 baseline, final packaged +Stage 2 is 2.276x and 2.919x faster, respectively. The small 128K difference +between the isolated 124.928 us/317.440 us values above and the packaged +125.952 us/316.416 us values is run-to-run variation. + +The final matched serving artifacts are under +`runs/bench/infllmv2/tmp/final_gqa_decode_20260813/`. They retain the same +workload ID, prompt-token hash, physical GPU UUID, and server arguments as the +historical comparisons. The sparse service was independently restarted and +measured twice; neither timed run enabled selector tracing. + +| Metric | Dense | Final sparse run 1 | Final sparse repeat | +| --- | ---: | ---: | ---: | +| Selector block hit rate | — | 57.5521% (884 / 1536) | Same bound diagnostic | +| Request throughput (req/s) | 0.092431 | 0.073803 (0.798x) | 0.074985 (0.811x) | +| Output throughput (token/s) | 11.831154 | 9.446807 (0.798x) | 9.598114 (0.811x) | +| Total-token throughput (token/s) | 3040.606665 | 2427.829371 (0.798x) | 2466.715343 (0.811x) | +| Mean / median / P99 TTFT (ms) | 7394.920 / 7395.463 / 7436.516 | 11214.213 / 11260.205 / 11902.204 | 10500.386 / 10554.044 / 10999.020 | +| Mean / median / P99 TPOT (ms) | 26.876 / 26.880 / 27.012 | 18.315 / 19.624 / 25.340 | 22.248 / 22.148 / 24.921 | +| Baseline / peak / peak-delta GPU memory (MiB) | 19493 / 24421 / 4928 | 19417 / 24345 / 4928 | 19417 / 24345 / 4928 | + +Both sparse runs reduce client-observed mean TPOT: by 31.85% (1.467x) and +17.22% (1.208x). This is the first matched 32K result in this implementation +where sparse TPOT is lower than dense TPOT. It is not a universal end-to-end +speedup: sparse TTFT is 42.0%-51.6% higher and the 128-token output throughput +is only 0.798x-0.811x dense. + +Streaming is also bursty in the final sparse runs: 45.08% and 37.20% of ITLs +are below 1 ms, versus 0% for dense. The benchmark's clean rate after excluding +those bursts is 37.21 token/s for dense and 30.04/28.26 token/s for sparse. +Consequently, the report treats the two lower TPOT measurements as reproducible +client-observed latency, reports their range, and does not reinterpret the +single best 1.467x value as stable kernel throughput. The direct CUDA-event and +Nsight results above remain the kernel-level evidence. + +This resolves why sparse is theoretically attractive without assuming it must +win every metric: it reads only 18.75% of the dense Stage 2 tokens at 32K, but +still pays selector, merge, launch, prefill/TTFT, and non-attention model costs. +The complete raw results and this interpretation are summarized in +`runs/bench/infllmv2/tmp/final_gqa_decode_20260813/summary.md`. + +### CUDA-restored baseline rerun (2026-08-16) + +The checkpoint sections above are retained as optimization history. The +current-source baseline at that checkpoint is +`runs/bench/infllmv2/service_gate_65k8_20260816/summary.md`; it was produced +after a fresh SM86 rebuild and does not reuse those historical conclusions. +The build output and loaded repository library are byte-identical, size +1,040,001,400 bytes and SHA-256 +`05e5b67e43d0da09bcbb884ae5b78fba23b8a3a99e31be69a3216c7a5d00e2c7`. +The focused current-package regression passed 24/24 tests, including +cross-page compressed-K, Stage 1 score/Top-K, paged Stage 2, and independent +NumPy/Paddle references. + +On the model-valid 65,344-token, batch-8 BF16 workload (`QH=32`, `KVH=2`, +`D=128`, 96 selected pages), CUDA-event medians are: + +| Implementation | Stage 1 + Top-K (us) | Stage 2 (us) | Sparse chain (us) | Dense baseline (us) | +| --- | ---: | ---: | ---: | ---: | +| FastDeploy | 349.184 | 318.464 | 650.240 | 785.408 | +| local `infllmv2_cuda_impl` | 318.464 | 208.896 | 573.440 | 1,563.648 | + +FastDeploy passes its own operator gate by 1.2079x. Its Stage 1 + Top-K, +Stage 2, and sparse chain are respectively 1.0965x, 1.5245x, and 1.1339x +slower than the local implementation. The dense columns are not comparable +across implementations because their dense references differ. The synthetic +page-aligned operator rate is 98,304 / 1,045,504 = 9.40255%. In the bound +service trace, partial causal pages make the actual rate 97,332 / 1,045,556 = +9.30911%, versus a 9.40208% full-page preset (-0.09296 percentage point). + +Nsight identifies Stage 2 main as the largest sparse kernel (286.205 us), +followed by fine block-score (249.816 us). A single 64-token shared-page Stage +2 experiment was rejected: although it improved 32K/batch-1, it regressed +128K/batch-4 Stage 2 by 17.1%, and Nsight attributed the regression to the +47.3 KiB shared-memory tile reducing occupancy. The retained 16-token tile +preserves GQA K/V sharing, split-KV, and online softmax. + +Dense and sparse services were then completely restarted and measured with +the same 65,344-input/128-output, eight-request, concurrency-8 workload. The +validated workload ID is +`4068415f92b013c216341f8461e29f42c979aa0aa76070b9677844817d68a1fd`. +Sparse did not beat dense end to end: output throughput was 5.439 versus 5.531 +token/s (0.983x), mean TTFT was 108,132 versus 105,124 ms (1.029x), and mean +TPOT was 628.660 versus 627.799 ms (1.001x). Selector micro-recall was 53.84% +(827/1,536). + +### Vectorized paged-K/V rerun (2026-08-17) + +The authoritative current-source evidence is +`runs/bench/infllmv2/service_gate_vector_load_20260817/summary.md`. The SM86 +package was rebuilt again. The build output and repository-loaded library are +byte-identical, size 1,040,001,400 bytes and SHA-256 +`1370f08391426352349d420b9268a769d8960615d32a717f1bb66a349dca59dc`. +The source header SHA-256 is +`883dbc1f5bfbfa8748bb83b3436766f7681d034a26de2e1811661aa738ccd6ae`. +The focused regression again passed 24/24 tests. + +This round changed only full-tile paged K/V movement in the production Stage +2 kernel: the scalar global-to-shared loop became aligned 16-byte `uint4` +loads/stores. Partial causal tiles keep the explicit scalar masking path. +Stage 1, CUB radix Top-K, the 16-token data layout, four-warp tensor-core +parallel structure, two-pages-per-split policy, online softmax, and split +combine remain unchanged. This differs from the local repository's 64-token +FlashAttention/CUTE tile only where paged-cache occupancy evidence requires it; +no simpler scalar attention algorithm was substituted. + +With the same 65,344-token, batch-8 BF16 workload and 20 warmups plus 100 CUDA +event measurements, the new medians are: + +| Implementation | Stage 1 + Top-K (us) | Stage 2 (us) | Sparse chain (us) | Dense baseline (us) | +| --- | ---: | ---: | ---: | ---: | +| FastDeploy | 329.728 | 163.840 | 507.424 | 784.384 | +| local `infllmv2_cuda_impl` | 304.112 | 206.848 | 547.840 | 1,553.408 | + +Relative to the 2026-08-16 FastDeploy baseline, Stage 2 is 1.9438x faster and +the sparse chain is 21.96% lower. FastDeploy now passes its own dense gate by +1.5458x. Against the freshly rerun local implementation, FastDeploy Stage 2 is +1.2625x faster and its sparse chain is 1.0796x faster, while its Stage 1 + +Top-K remains 1.0842x slower. Dense references remain framework-specific and +are not compared across implementations. + +Nsight Systems measures the Stage 2 main kernel at 145.918 us, down from +286.205 us. Fine block-score is now the largest sparse kernel at 255.564 us; +the next optimization round must therefore target Stage 1 rather than tune +Stage 2 block counts. The operator effective rate is 98,304 / 1,045,504 = +9.40255%. The current bound service trace is 97,332 / 1,045,556 = 9.30911%, +versus a 9.40208% full-page preset (-0.09296 percentage point). + +Dense and untraced sparse services were fully restarted with the same validated +workload ID. The complete service still did not beat dense: sparse output +throughput was 4.832 versus 4.915 token/s (0.983x), mean TTFT was 133,586 versus +131,160 ms (1.018x), and mean TPOT was 614.727 versus 605.495 ms (1.015x). +Selector micro-recall was 53.39% (820/1,536). The operator gain is real, but it +does not justify a universal end-to-end speedup claim for this long-prefill, +128-token-output workload. + +## W4A16 attention and serving E2E record (2026-08-25) + +This section uses one convention throughout: latency speedup is +`FlashAttention latency / InfLLM-v2 latency`, while throughput speedup is +`InfLLM-v2 throughput / FlashAttention throughput`. Values above 1 mean that +InfLLM-v2 is faster. Attention-only measurements must not be confused with +service E2E, which also includes W4A16 Linear layers, MLPs, normalization, +sampling, scheduling, and API overhead. + +### Current-source standalone attention measurements + +The isolated setup used one RTX A6000 (SM86), the MiniCPM4.1 production shape +`QH=32`, `KVH=2`, `D=128`, `block_size=64`, batch 1, and BF16 Q/K/V plus BF16 +KV cache. Prefill used 20 warmups and 100 CUDA-event samples; decode used 50 +warmups and 200 samples. + +| Phase | Context | FlashAttention median | InfLLM-v2 complete-chain median | Speedup | +| --- | ---: | ---: | ---: | ---: | +| Prefill attention | 16,384 | 20.824 ms | 18.514 ms | 1.125x | +| Prefill attention | 32,768 | 88.638 ms | 45.810 ms | 1.935x | +| Decode attention | 32,768 | 62.464 us | 181.248 us | 0.345x | + +The 32K decode chain comprises a 16.384 us summary update, 95.232 us Stage 1, +and 58.368 us Stage 2. Stage 2 alone is 1.070x faster than dense attention, but +the complete chain is only 0.345x, or 2.902x slower. A Stage-2-only number is +therefore not a decode-attention speedup. See the raw +[prefill JSON](../../runs/bench/infllmv2/w4a16_sm86_tp1_20260825/attention_prefill_16k_32k_b1.json) +and [decode JSON](../../runs/bench/infllmv2/w4a16_sm86_tp1_20260825/attention_decode_32k_b1.json). + +### W4A16 serving definition and current status + +W4A16 means INT4 Linear weights with BF16 activations. Attention Q/K/V and KV +cache remain BF16. The only intended comparison variable was the attention +backend; both services otherwise used TP1, `max_model_len=65536`, +`block_size=64`, `max_num_seqs=1`, `max_num_batched_tokens=32768`, with prefix +caching, chunked prefill, and CUDA Graph disabled. The workload was four exact +32,768-input/128-output requests at concurrency 1 and seed 2026. + +The required ordering was enforced: load the original BF16 checkpoint, perform +online INT4 weight quantization during startup with `--quantization wint4`, +mark workers ready, then warm up and time requests. Request E2E therefore +excludes one-time startup quantization. Quantized weights have a theoretical +4x bit-width reduction; actual model or GPU-memory reduction also includes +scales, unquantized parameters, and the BF16 KV cache. + +The current FlashAttention baseline succeeded with mean TTFT 7,802.485 ms, +TPOT 14.034 ms, and output throughput 13.345 token/s. InfLLM-v2 activated +sparse prefill during the untimed 32K warmup, then exited with +`CUBLAS_STATUS_INVALID_VALUE` in the LM head; disabling overlap scheduling +produced the same failure. Consequently there is no honest current-source +InfLLM-v2 W4A16 serving speedup. A failed run is neither 0x nor a reason to +substitute an attention-only ratio. See the successful +[FlashAttention result](../../runs/bench/infllmv2/w4a16_sm86_tp1_20260825/serving_flash_attn_32k_128_b1.json) +and the [InfLLM-v2 failure record](../../runs/bench/infllmv2/w4a16_sm86_tp1_20260825/serving_infllmv2_failure.json). + +### Retained successful W4A16 service pair (2026-08-19, historical checkpoint) + +The repository also retains a successful paired run that permits a reproducible +serving calculation. Both backends used the same offline-prequantized W4A16 +checkpoint, one RTX A6000, two 65,344-input/16-output requests, and concurrency +1. Both input payloads have SHA-256 +`1100aec67962ce4cf3cf5e9bc18ab7c1bcc44ff06bbf71d80b058322b4a02aa8`. +This is not a 2026-08-25 current-source result because the current loader no +longer supports that prequantized checkpoint path. + +| Serving metric | FlashAttention | InfLLM-v2 | InfLLM-v2 speedup | +| --- | ---: | ---: | ---: | +| Prefill proxy: mean TTFT | 21,628.915 ms | 12,950.952 ms | 1.670x | +| Decode proxy: mean TPOT | 18.143 ms | 35.478 ms | 0.511x | +| Derived mean E2E/request | 21,901.054 ms | 13,483.121 ms | 1.624x | +| Measured total duration | 43.820 s | 26.983 s | 1.624x | +| Output throughput | 0.730 token/s | 1.186 token/s | 1.624x | +| Clean decode throughput | 53.332 token/s | 28.186 token/s | 0.529x | + +Here `mean E2E = mean TTFT + (16 - 1) * mean TPOT`. The duration and output +throughput ratios agree because the token totals are identical. Sparse prefill +substantially improves TTFT, while decode is about 1.956x slower. This +short-output workload remains prefill-dominated, yielding about 1.624x overall +E2E speedup. The raw records are +[FlashAttention](../../runs/bench/infllmv2/w4a16_64k_b1_gpu7_20260819/flash_attn_64k_single_prefill.json) +and [InfLLM-v2](../../runs/bench/infllmv2/w4a16_64k_b1_gpu7_20260819/infllmv2_64k_single_prefill.json); +the consolidated calculation is in [summary.json](../../runs/bench/infllmv2/w4a16_sm86_tp1_20260825/summary.json). + +## Acceptance matrix + +| Item | Code/evidence | Status when this document was written | +| --- | --- | --- | +| INF-1: per-request/per-head Stage 1 | `infllmv2_update_compressed_k`, `infllmv2_select_blocks`; backend/operator tests | Current package passed 24/24 focused cases: backend 15/15 and packaged GPU 9/9 | +| INF-2: post-RoPE + paged-cache Stage 2 | `decoder_write_cache_with_rope` before `infllmv2_attention_forward`; ordering and paged-mapping tests | Passed: writer order, already-written cache, cross-page summaries, and paged mapping are covered | +| P0: GQA-tiled Stage 2 | Production-shape CUDA path and NumPy-oracle operator test | Passed; aligned paged K/V movement reduced current 65K/8 Stage 2 to 163.840 us versus 206.848 us locally | +| P2: persistent Stage 2 buffers | Backend workspace reuse and four-output in-place alias tests | Passed; no Stage 2 hot-path `empty`/`full`/`zeros` allocation | +| Stage 1-1: GQA-tiled coarse LSE | Split partial/merge kernels, persistent workspaces, semantic-reference GPU test, Nsight traces | Passed; 14.558x at 32K/1 and 21.151x at 128K/4 | +| Stage 1-2: exact radix TopK | Composite-key CUB selector and 2,048-candidate tie/padding GPU tests | Passed; 2.478x at 32K/1 and 22.545x at 128K/4 | +| Stage 1-3: GQA-tiled block-score | Eight-block K tiles shared across 16 query heads; production GPU oracle tests and Nsight traces | Passed; block-score is 4.382x faster at 32K/1 and 9.391x faster at 128K/4 than the preceding checkpoint | +| P3: calibrated long-context comparison | Current CUDA-event gate, same-shape local-repository rerun, Nsight, and fully restarted 65K/8 service report | Operator gate passed at 1.5458x; current complete service remained at 0.983x dense output throughput, so no universal service-speedup claim is made | +| Implementation documentation | This page, root `rfc.md`, and the current rerun summary | Provided with traceable current-source results; the 2026-08-17 section supersedes earlier performance checkpoints | + +## References + +- [InfLLM-V2 paper](https://arxiv.org/abs/2509.24663) +- [OpenBMB InfLLM-V2 CUDA implementation](https://github.com/OpenBMB/infllmv2_cuda_impl) +- [MiniCPM4.1-8B](https://huggingface.co/openbmb/MiniCPM4.1-8B) +- [AngelSlim/Hy3-GGUF](https://huggingface.co/AngelSlim/Hy3-GGUF) +- [llama.cpp PR #25395](https://github.com/ggml-org/llama.cpp/pull/25395) + +Hy3-GGUF and llama.cpp PR #25395 are user-specified low-bit/Hy3 engineering references. They are not sources for this InfLLM-V2 selection algorithm or its performance numbers. diff --git a/docs/features/thinking_budget.md b/docs/features/thinking_budget.md index c236cb4fcb4..4e06d66fc04 100644 --- a/docs/features/thinking_budget.md +++ b/docs/features/thinking_budget.md @@ -14,14 +14,29 @@ segment. When the budget is reached, it terminates thinking by forcing ` ## How It Works 1. **Request-side precompute (DataProcessor)**: when a request includes `thinking_budget`, the prompt token ids are scanned to determine whether thinking has started, whether it already ended, and how many tokens are already inside the thinking section. + Single-token markers can be precomputed by the `DataProcessor`; multi-token markers are replayed by the logits processor so partial marker prefixes are preserved. 2. **Per-step update**: during decoding, the processor tracks `last_token_id` and `tokens_after_start`. 3. **Budget enforcement**: once the budget is reached, it forces `` directly. If `think_stop_sentence` is configured, it forces that sentence first and then ``. ## Requirements -- The model must provide valid token ids for `think_start_id` and `think_end_id` (via `ModelConfig`). -- If either of these ids is invalid, the processor is disabled and `thinking_budget` will not take effect. +- The model must provide either a valid `think_start_id`/`think_end_id` pair or a complete `think_token_sequences` configuration. +- If neither mode is configured, the generic processor is disabled. A model that explicitly declares multi-token marker support fails startup when its tokenizer markers cannot be derived. + +### Multi-token thinking markers (`think_token_sequences`) + +Some tokenizers (e.g. MiniCPM4.1) do not have single ``/`` vocab entries: the markers +tokenize to multiple token ids that vary with context (word-start vs. mid-text variants). In that case +the model exposes a tokenizer-sequence builder hook. The engine invokes that hook at startup and passes +standalone/contextual markers plus the complete forced end sequence through `ModelConfig.think_token_sequences` +(structure: `{"start": [[...]], "end": [[...]], "forced_end": [...]}`): + +- Thinking segments are detected by token-sequence matching; sequence prefixes do not consume budget. +- Once the budget is reached, the required `forced_end` sequence is emitted token by token and each emitted + token is verified on the following decode step. +- Sequence derivation is selected by the explicit model hook, not by the broad `REASONING` model category; + builder and configuration errors propagate and stop startup. ## Request Parameters diff --git a/docs/get_started/minicpm41.md b/docs/get_started/minicpm41.md new file mode 100644 index 00000000000..dbc80b955fe --- /dev/null +++ b/docs/get_started/minicpm41.md @@ -0,0 +1,276 @@ +[简体中文](../zh/get_started/minicpm41.md) + +# Deploy the MiniCPM4.1-8B Model + +This document explains how to deploy MiniCPM4.1-8B in BF16 or with online WINT4/WINT8 quantization. Before starting the deployment, ensure that your hardware environment meets the following requirements: + +- GPU Driver >= 535 +- CUDA >= 12.3 +- CUDNN >= 9.5 +- Linux X86_64 +- Python >= 3.10 +- 1 x 48 GB NVIDIA GPU + +For FastDeploy installation instructions, refer to the [Installation Guide](./installation/README.md). + +## Minimal Validation Setup + +Enter the FastDeploy repository and specify only the local model directory and the GPU used for testing. The scripts automatically use the project `.venv` and configure its CUDA runtime libraries: + +```shell +cd /path/to/FastDeploy +export MODEL_PATH=/path/to/MiniCPM4.1-8B +export CUDA_VISIBLE_DEVICES=0 +``` + +## Prepare the Model + +### 1. Manual Download (Optional) + +MiniCPM4.1-8B uses the Hugging Face Torch safetensors format. Run the following commands to download the model: + +```shell +export MODEL_PATH=/path/to/MiniCPM4.1-8B +hf download openbmb/MiniCPM4.1-8B --local-dir "${MODEL_PATH}" +``` + +The `pad_token_id` in the model configuration is 2. Before deployment, ensure that `pad_token` in `tokenizer_config.json` is set to ``, which corresponds to token ID 2. Run the following command to complete the configuration: + +```shell +MODEL_PATH="${MODEL_PATH}" python - <<'PY' +import json +import os +from pathlib import Path + +config_path = Path(os.environ["MODEL_PATH"]) / "tokenizer_config.json" +config = json.loads(config_path.read_text(encoding="utf-8")) +config["pad_token"] = "" +config_path.write_text( + json.dumps(config, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", +) +PY +``` + +### 2. Automatic Download + +Set the model name to `openbmb/MiniCPM4.1-8B`: + +```shell +export MODEL_PATH=openbmb/MiniCPM4.1-8B +``` + +>💡 **Note**: If the path specified by `--model` does not exist in the current directory, FastDeploy queries AIStudio for a preset model with that name. When found, the model is downloaded automatically with resumable transfers. AIStudio is the default download source; use `FD_MODEL_SOURCE` and `FD_MODEL_CACHE` to configure the source and cache directory. For details, see [Model Download](../supported_models.md). If the model has already been downloaded, you can set `MODEL_PATH` to its local directory instead. + +## Build and Minimal Regression + +### Build the CUDA Operators + +Build the CUDA operators and verify the InfLLM-V2 symbols: + +```shell +bash tests/benchmarks/test_minicpm41.sh build +``` + +The script invokes `build.sh`. Set `FD_BUILDING_ARCS` when targeting a different GPU architecture. A successful run prints `PASS build` and verifies that `fastdeploy/model_executor/ops/gpu/fastdeploy_ops/fastdeploy_ops_pd_.so` contains these symbols: + +- `infllmv2_update_compressed_k` +- `infllmv2_select_blocks` +- `infllmv2_attention_forward` + +### MiniCPM4.1, Thinking, and Quantization Tests + +Test model registration, weight mapping, multi-token thinking, mixed thinking modes, WINT4/WINT8 online quantization, and the InfLLM-V2 backend: + +```shell +bash tests/benchmarks/test_minicpm41.sh unit +``` + +Run correctness tests for the compiled CUDA operators: + +```shell +bash tests/benchmarks/test_minicpm41.sh operators +``` + +A zero exit status indicates success. + +### End-to-End Serving Tests + +Run BF16, WINT4, WINT8, and InfLLM-V2 validation sequentially with one command: + +```shell +bash tests/benchmarks/test_minicpm41.sh e2e +``` + +The E2E runner automatically selects temporary ports, starts and cleans up each server, and runs five cases from `test_minicpm41_serving.py` in every mode. Each mode passes when it reports `5 passed`. + +To run everything from compilation through E2E validation: + +```shell +bash tests/benchmarks/test_minicpm41.sh +``` + +## Start the Service + +>💡 **Note**: The following command serves the BF16 model on a single GPU with prefix caching disabled. + +Run the following command to start the service. For details about the startup options, refer to the [Parameter Guide](../parameters.md). + +```shell +export CUDA_VISIBLE_DEVICES=0 +export FD_ATTENTION_BACKEND=FLASH_ATTN + +python -m fastdeploy.entrypoints.openai.api_server \ + --model "${MODEL_PATH}" \ + --served-model-name MiniCPM4.1-8B \ + --port 8180 --engine-worker-queue-port 8182 \ + --cache-queue-port 8183 --metrics-port 8181 \ + --tensor-parallel-size 1 \ + --max-model-len 8192 \ + --max-num-seqs 1 \ + --max-num-batched-tokens 128 \ + --no-enable-prefix-caching +``` + +### Online WINT4/WINT8 + +Use the same original BF16 checkpoint and add `--quantization wint4` or `--quantization wint8`. FastDeploy quantizes each Linear weight in memory after loading it; no converted checkpoint is needed or supported by this MiniCPM4.1 path. The other service options are identical to the BF16 command above. + +After completing the minimal validation setup, start WINT4: + +```shell +bash scripts/run_minicpm41_wint_server.sh wint4 "${MODEL_PATH}" +``` + +Replace `wint4` with `wint8` for online INT8 weight quantization. The script requires `MODEL_PATH` to point to a local model directory and automatically uses the project `.venv`, its CUDA runtime libraries, and safe single-GPU defaults. Additional server options can be appended to the command. If the ports are already in use, assign an unused set of API, metrics, engine queue, and cache queue ports. See [Online Quantization](../quantization/online_quantization.md) for the shared WINT behavior. + +## End-to-End Performance and Compression Results + +The following results were measured with the original MiniCPM4.1-8B BF16 checkpoint and one NVIDIA RTX A6000 (SM86, TP=1). All modes used `max_model_len=1024`, `max_num_seqs=1`, `max_num_batched_tokens=256`, prefix caching disabled, and 32 fixed GPU KV-cache blocks. The serving benchmark sent 16 random requests at concurrency 1 with nominal 128-token inputs and 64-token outputs; the identical completed workload in every mode contained 2,895 input and 1,010 output tokens. + +| Mode | Output tok/s | E2E speedup | Mean TTFT | Mean TPOT | TPOT speedup | Resident GPU memory | Memory compression | Memory reduction | Estimated parameter storage | Parameter compression | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| BF16 | 37.35 | 1.00x | 136.98 ms | 24.45 ms | 1.00x | 16,858 MiB (16.46 GiB) | 1.00x | 0.00% | 15.25 GiB | 1.00x | +| Online WINT4 (INT4 weight) | 96.72 | **2.59x** | 146.50 ms | 7.86 ms | 3.11x | 6,514 MiB (6.36 GiB) | **2.59x** | 61.36% | 4.66 GiB | 3.28x | +| Online WINT8 (INT8 weight) | 64.48 | **1.73x** | 140.23 ms | 13.24 ms | 1.85x | 10,416 MiB (10.17 GiB) | **1.62x** | 38.21% | 8.19 GiB | 1.86x | +| W4AFP8 | N/A | N/A | N/A | N/A | N/A | N/A | N/A | N/A | N/A | N/A | + +The end-to-end speedup is `quantized output throughput / BF16 output throughput`. Because all three runs produced the same 1,010 output tokens, it is also the inverse benchmark-duration ratio. Memory compression is `BF16 resident GPU memory / quantized resident GPU memory`; the 15 MiB idle GPU value was removed, and all modes reserved the same 32 KV-cache blocks. Estimated parameter storage is calculated from safetensors metadata and includes BF16 weights that are not quantized plus persistent BF16 scales. It is therefore different from resident GPU memory, which also includes runtime allocations and the KV cache. Online quantization slightly increased worker startup time from 25.41 seconds for BF16 to 27.41 seconds for WINT4 and 28.09 seconds for WINT8. + +W4AFP8 does not have a valid result on this host. For dense Linear layers, the `w4afp8` CLI selection currently resolves to block-wise FP8, while the actual W4AFP8 implementation applies to `FusedMoE`; MiniCPM4.1-8B is dense. An attempted launch on SM86 failed during CUDA graph capture when DeepGEMM tried to build an SM90a FP8 kernel. Consequently, neither an end-to-end speedup nor an actual memory-compression ratio is reported. A nominal 4-bit weight payload may be described as roughly 4x smaller than BF16, but that theoretical payload ratio is not a runnable MiniCPM4.1 W4AFP8 result. + +Reproduce the serving workload after starting each server with the settings above: + +```shell +python benchmarks/benchmark_serving.py \ + --backend openai-chat \ + --base-url http://127.0.0.1:8180 \ + --endpoint /v1/chat/completions \ + --model "${MODEL_PATH}" --tokenizer "${MODEL_PATH}" \ + --dataset-name random --num-prompts 16 \ + --random-input-len 128 --random-output-len 64 \ + --max-concurrency 1 --seed 0 --save-result +``` + +## Send Requests to the Service + +After you run the service startup command, the following terminal output indicates that the service has started successfully. + +```shell +INFO api_server.py[line:1030] Launching metrics service at http://0.0.0.0:8181/metrics +INFO api_server.py[line:1033] Launching chat completion service at http://0.0.0.0:8180/v1/chat/completions +INFO api_server.py[line:1034] Launching completion service at http://0.0.0.0:8180/v1/completions +[INFO] Starting gunicorn 26.0.0 +[INFO] Listening at: http://0.0.0.0:8180 +[INFO] Application startup complete. +``` + +FastDeploy provides a health check endpoint for checking the service status. If the following command returns `HTTP/1.1 200 OK`, the service has started successfully. + +```shell +curl -i http://0.0.0.0:8180/health +``` + +Send a request using the following command. `enable_thinking=false` disables thinking mode. + +```shell +curl -X POST "http://0.0.0.0:8180/v1/chat/completions" \ +-H "Content-Type: application/json" \ +-d '{ + "model": "MiniCPM4.1-8B", + "messages": [ + {"role": "user", "content": "Rewrite the poem Quiet Night Thought by Li Bai as a modern poem."} + ], + "temperature": 0, + "top_p": 1, + "max_tokens": 64, + "stream": false, + "chat_template_kwargs": {"enable_thinking": false} +}' | jq --indent 4 . +``` + +The response is as follows: + +```json + +{ + "object": "chat.completion", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Certainly. Here is a modern-poetry adaptation of Li Bai's \"Quiet Night Thought\":\n\n**Night Thoughts**\n\nOutside the window, moonlight\npours down like flowing water\n\nA vast expanse of white\ncovers the cold window lattice\n\nA distant home\nflickers deep within memory\n\nA little bed\n" + }, + "logprobs": null, + "draft_logprobs": null, + "prompt_logprobs": null, + "finish_reason": "length", + "speculate_metrics": null + } + ], + "usage": { + "prompt_tokens": 29, + "total_tokens": 93, + "completion_tokens": 64, + "prompt_tokens_details": { + "cached_tokens": 0, + "image_tokens": 0, + "video_tokens": 0 + }, + "completion_tokens_details": { + "reasoning_tokens": 0, + "image_tokens": 0 + } + } +} +``` + +FastDeploy's service API is compatible with the OpenAI protocol. You can send a request using the following Python code: + +```python +import openai + +host = "0.0.0.0" +port = "8180" +client = openai.Client(base_url=f"http://{host}:{port}/v1", api_key="null") + +response = client.chat.completions.create( + model="MiniCPM4.1-8B", + messages=[ + {"role": "system", "content": "I'm a helpful AI assistant."}, + {"role": "user", "content": "Rewrite the poem Quiet Night Thought by Li Bai as a modern poem."}, + ], + temperature=0, + top_p=1, + max_tokens=64, + stream=True, + extra_body={"chat_template_kwargs": {"enable_thinking": False}}, +) +for chunk in response: + if chunk.choices[0].delta.content: + print(chunk.choices[0].delta.content, end="") +print("\n") +``` + +MiniCPM4.1 supports both thinking and non-thinking modes. Set `chat_template_kwargs.enable_thinking` to `true` in the request to enable thinking mode. You can also use `reasoning_max_tokens` to limit the number of thinking tokens used for the request. diff --git a/docs/supported_models.md b/docs/supported_models.md index b0684affc11..d73cbad1a60 100644 --- a/docs/supported_models.md +++ b/docs/supported_models.md @@ -40,6 +40,7 @@ These models accept text input. |⭐DEEPSEEK|BF16/WINT4|unsloth/DeepSeek-V3.1-BF16;
unsloth/DeepSeek-V3-0324-BF16;
unsloth/DeepSeek-R1-BF16, etc.| |⭐GPT-OSS|BF16/WINT8|unsloth/gpt-oss-20b-BF16, etc.| |⭐GLM-4.5/4.6|BF16/wfp8afp8|zai-org/GLM-4.5-Air;
zai-org/GLM-4.6
 [最佳实践](./best_practices/GLM-4-MoE-Text.md) etc.| +|⭐MiniCPM4.1|BF16/WINT4/WINT8|openbmb/MiniCPM4.1-8B
 [测试教程](./get_started/minicpm41.md)| ## Multimodal Language Models diff --git a/docs/zh/features/infllmv2_attention.md b/docs/zh/features/infllmv2_attention.md new file mode 100644 index 00000000000..ebc32a893ac --- /dev/null +++ b/docs/zh/features/infllmv2_attention.md @@ -0,0 +1,768 @@ +[English](../../features/infllmv2_attention.md) + +# InfLLM-V2 稀疏注意力 + +FastDeploy 的 `INFLLMV2_ATTN` 后端为 MiniCPM4.1 的 decode 和长上下文 prefill 提供两阶段稀疏注意力：Stage 1 按请求和 KV head 动态选择逻辑 cache 块，Stage 2 通过 `block_tables` 读取已写入的 paged K/V cache。 + +初始 batch-1 prefill 在 `dense_len` 之前使用 dense FlashAttention，之后每个 128-token query tile 共享一次 Stage 1 选块。每个稀疏 tile 被拆成“完全可见的历史选中块”和“当前一至两个块”：历史分区运行 non-causal FlashAttention，当前分区运行 causal FlashAttention，再用两边精确的 log-sum-exp 合并输出。这样即使 Paddle FlashAttention V2 会把较短 causal query 左上对齐，也能保持正确的因果位置。短请求、mixed batch 和共享 chunked-prefill 路径继续使用 dense prefill，同时仍建立 decode 所需的双尺度语义摘要。 + +## 实现概览 + +decode 数据流如下： + +```text +raw fused QKV + -> decoder_write_cache_with_rope + -> post-RoPE / post-QK-norm Q + -> 当前 K/V 写入 paged cache + -> infllmv2_update_compressed_k + -> infllmv2_select_blocks + -> infllmv2_attention_forward + -> block_tables[请求, 逻辑块] -> physical page + -> causal sparse softmax +``` + +Stage 1 和 Stage 2 都使用 writer 产生的 post-RoPE query，Stage 2 读取的 K/V 包含当前 decode token。原始 fused QKV 不会被直接送入 sparse attention。 + +sparse prefill 数据流如下： + +```text +raw fused QKV + -> gqa_rope_write_cache + -> post-RoPE / post-QK-norm Q/K/V 和 paged K/V cache + -> infllmv2_update_compressed_k + -> [0, dense_len) 使用 dense FlashAttention + -> 每个 128-token query tile 调用一次 infllmv2_select_blocks + -> gather 选中的 paged K/V + -> 历史选中页运行 non-causal FlashAttention + -> 当前一至两个页运行 causal FlashAttention + -> 使用精确 LSE 加权合并输出 +``` + +默认每次处理 4,096 个 token，即 32 个 query tile，以限制临时 gather K/V 的显存；最后不足 128 token 的 tile 也受支持，prompt 长度不必按 block 对齐。 + +对每个 query token，请求与 causal 位置由运行时元数据计算： + +```text +request = batch_id_per_token[token] +local_offset = token - cu_seqlens_q[request] +position = seq_lens_decoder[request] + local_offset +``` + +因此 continuous batching 中的不同请求、KV head 和 paged-cache block table 保持隔离；`batch_id_per_token == -1` 的 padding token 不会访问 cache。 + +### 双尺度 Stage 1 + +Stage 1 使用两种 K 语义窗口： + +| 尺度 | 默认窗口 | 默认步长 | 用途 | +| --- | ---: | ---: | --- | +| fine | 32 | 16 | 为逻辑 KV 块计算精细相关度 | +| coarse | 128 | 64 | 估计每个 query head 的 log-sum-exp 归一化项 | + +每个窗口保存 paged K cache 中 K 向量的均值。完整窗口归属于包含其最后一个 token 的 physical page，所以 fine/coarse 窗口可以跨 page。对默认 `block_size=64`，每个 physical block 和 KV head 有 4 个 fine slot 和 1 个 coarse slot。 + +设 `G(g)` 是 KV head `g` 对应的 GQA query-head 组，`d` 是 head dimension。打分过程为： + +```text +LSE[h] = logsumexp_j(dot(Q[h], Kbar_coarse[j]) / sqrt(d)) + +semantic_score[i, g] = sum(h in G(g)) exp( + dot(Q[h], Kbar_fine[i]) / sqrt(d) - LSE[h] +) + +block_score[b, g] = max(semantic_score[i, g] + for fine window i overlapping block b) +``` + +对 MiniCPM4.1 生产形状（`QH=32`、`KVH=2`、`D=128`、`block_size=64`），coarse 归一化使用参考 OpenBMB CUDA 实现分布式 softmax 结构改写的 GQA-tiled +split kernel。一个 CTA 负责 `(query token, KV head, coarse split)`；16 个 half-warp 将 16 个 Q head 常驻寄存器，协作加载一份对齐的 `float4` K tile（最多16 个 coarse windows），并在整个 GQA 组内复用。每个 split 的 online`(max, sum)` 分布在各 lane，随后由“一 warp/一 query head”的 kernel 做 rescale归并。split partial 是调用方提供的持久化 workspace，因此该路径既没有thread-0 串行窗口循环，也没有逐次调用分配。 + +生产形状的 fine block-score 路径一次 tile 8 个相邻候选块。一个 CTA 负责 `(query token, KV head, 8-block tile)`；16 个 half-warp 将 GQA query head 常驻寄存器，最多 33 个重叠 fine K row 只通过对齐 `float4` 搬运一次并由整个 GQA 组共享，随后同一组 head contribution 同时为 8 个候选块做 pooling。backend 会跨 decode step 复用 score workspace，因此 kernel 会重写所有 score slot， +包括不足 8 块的 table 尾部。 + +初始 `init_blocks` 和最近 `window_size` 对应的局部块会被置为正无穷分并参与同一次 top-k，其余候选块按 `block_score` 排序。`local_blocks = window_size / block_size` 且包含当前块，进入稀疏区后的选中预算为 `topk + local_blocks`；因此 `topk` 不是强制块之外的额外预算。默认配置会选中 96 块，其中包含 33 个强制块（1 个初始块和含当前块在内的 32 个局部块）与最多 63 个普通动态块。输出 buffer 为了容纳短上下文选全块，其容量为： + +```text +selected_capacity = max( + topk + window_size / block_size, + ceil(dense_len / block_size), +) +``` + +稀疏区选块使用精确的 64-bit 复合 radix key：分数降序、逻辑块 ID 升序。 +256-thread CUB block radix sort 按候选规模为每线程分派 1、2、4 或 8 个元素， +单 CTA 可覆盖最多 2,048 个逻辑块（`block_size=64` 时对应 128K token）。它替换 +了旧的全候选两两排名，同时保持 tie 时的确定性。短于 `dense_len` 的请求直接 +输出所有可见块，不进入排序；超过 2,048 个候选块的模型配置会显式失败。 + +Stage 1 在 GPU 上输出： + +| 张量 | 形状 | 说明 | +| --- | --- | --- | +| `topk_indices` | `[tokens, kv_heads, selected_capacity]` | 每请求/query token/KV head 的逻辑块 id，未用位为 `-1` | +| `block_scores` | `[tokens, kv_heads, max_blocks_per_seq]` | 动态块分数 | +| `selected_counts` | `[tokens, kv_heads]` | 实际选中块数 | + +`topk_indices` 的有效前缀必须是严格递增、无重复且在当前请求可见范围内的逻辑块 ID，其余 slot 必须全为 `-1`。Stage 2 依据第一个 `-1` 判定前缀长度；这是 custom op 的输入前置条件，正常请求中由 Stage 1 保证。runtime 序列长度、`cu_seqlens_q`、token/request 映射及 `block_tables` 页所有权也必须由 FastDeploy 调度器保证自洽；不同活跃请求不得共享可写 physical page。Q/K/V 输入必须为有限值。 + +MiniCPM4.1 是 GQA 模型，因此元数据按 KV head 生成，不会为同一 GQA 组内的每个 query head 复制一份块列表。Stage 2 使用该 KV head 的块列表服务同组 query heads。 + +### Paged Stage 2 + +Stage 2 不假设 logical block 在显存中连续，而是对每个选中块执行： + +```text +physical_block = block_tables[request, logical_block] +``` + +随后仅对该 physical page 中不超过当前请求 `position` 的 token 计算 scaled dot-product attention。`topk_indices` 必须是逻辑块 id，不得传 physical page id。 + +MiniCPM4.1 的生产形状为 32 个 query heads、2 个 KV heads、`head_dim=128`、`block_size=64`，Stage 2 对此使用 GQA-tiled FlashDecoding 内核。一个 128-thread CTA 负责 `(query token, KV head, KV split)`，一次覆盖共享该 KV head 的全部 16 个 query heads。内核保留 16-token K/V shared-memory tile：warp 0 计算 16x16 QK tensor-core tile，thread 0 至 15 更新各 head 的 online `(max, sum)`，四个 warp 共同计算 probability-times-V tensor-core tiles，并更新 shared FP32 accumulator。完整 16-token K/V tile 从 paged cache 通过对齐的 16-byte `uint4` transaction 搬入 shared memory；因果边界的非完整 tile 显式走带 mask 的标量路径。独立 combine kernel 归并各 split 的 `(accumulator, max, sum)` partial。 + +每个 KV split 包含两个选中 page，在 batch 1 时仍可提供足够 CTA，同时不会恢复旧实现的 16 倍 K/V 全局显存读取。本地 CUDA 子仓库使用 FlashAttention/CUTE 64-token K/V tile；FastDeploy 保留较小 tile，是因为直接移植 64-token paged tile 需要 47.3 KiB shared memory，并使 128K/batch 4 Stage 2 回退 17.1%。当前 16-token paged tile 约使用 22 KiB，同时保留 GQA 复用、split-KV 与 online softmax。非生产形状走通用正确性路径，不代表 MiniCPM4.1 的性能路径。 + +每层 backend 持久化持有 Stage 1 元数据、最终 attention 输出以及 FP32 split `(accumulator, max, sum)` tensor；它们都作为 in-place 输入传给 custom op，并在形状不变的 decode step 间复用。Stage 2 CUDA launcher 的热路径不再调用 `paddle::empty`、`paddle::full` 或 `paddle::zeros`。 + +## 配置与启用 + +使用环境变量显式选择后端： + +```shell +export FD_ATTENTION_BACKEND=INFLLMV2_ATTN +``` + +参数优先从模型 `config.json` 的 `sparse_config` 读取，其次为模型顶层同名字段，最后使用 MiniCPM4.1 默认值。如需要覆盖，可在模型配置中添加： + +```json +{ + "sparse_config": { + "kernel_size": 32, + "kernel_stride": 16, + "topk": 64, + "dense_len": 8192, + "init_blocks": 1, + "window_size": 2048, + "sparse_prefill": true, + "prefill_query_chunk_size": 4096 + } +} +``` + +| 参数 | 默认值 | 约束 | +| --- | ---: | --- | +| `block_size` | 64 | 启动参数 `--block-size`；必须同时被 `kernel_stride` 和 `4 * kernel_stride` 整除 | +| `kernel_size` | 32 | 正整数 | +| `kernel_stride` | 16 | 正整数 | +| `topk` | 64 | 正整数 | +| `dense_len` | 8192 | 至少为 `4 * kernel_size` | +| `init_blocks` | 1 | 非负整数 | +| `window_size` | 2048 | `block_size` 的非负整数倍 | +| `sparse_prefill` | `true` | 是否启用模型专用的初始 prefill 稀疏路径 | +| `prefill_query_chunk_size` | 4096 | 128 的正整数倍 | + +`init_blocks` 必须小于 `topk`，以保证初始块和包含当前块的局部集合都能放入输出容量。 + +32K 长上下文单卡服务示例： + +```shell +export CUDA_VISIBLE_DEVICES=0 +export FD_ATTENTION_BACKEND=INFLLMV2_ATTN +export MODEL_PATH=/path/to/MiniCPM4.1-8B + +.venv/bin/python -m fastdeploy.entrypoints.openai.api_server \ + --model "${MODEL_PATH}" \ + --served-model-name MiniCPM4.1-8B \ + --port 8180 --metrics-port 8181 \ + --engine-worker-queue-port 8182 --cache-queue-port 8183 \ + --tensor-parallel-size 1 \ + --max-model-len 65536 \ + --block-size 64 \ + --max-num-seqs 1 \ + --max-num-batched-tokens 32768 \ + --gpu-memory-utilization 0.5 \ + --no-enable-prefix-caching \ + --graph-optimization-config '{"use_cudagraph": false}' +``` + +## 限制与显存费用 + +- 只支持 compute capability 8.0 及以上（SM80+）的 NVIDIA CUDA GPU，以及 FP32、FP16 和 BF16。 +- paged K/V cache 必须为 rank 4 且未量化；量化 KV cache 会显式失败。 +- query heads 数必须被 KV heads 数整除，`head_dim` 必须在 `[1, 256]` 内。 +- tensor parallel 不得复制 KV head：`tensor_parallel_size` 必须不大于每层的全局 KV-head 数。当前尚未实现跨副本的 Stage 1 分数归并，因此不支持的配置会显式失败。 +- sparse decode 需要同一 paged K cache 先经过 prefill 建立语义摘要。当前 P/D 分离不会传输这些摘要，因此不得在缺失摘要时进入 sparse decode。 +- 更换或 reset paged K cache 时必须一起 reset 语义摘要。 +- 语义摘要尚未纳入 CUDA Graph 的 cache replacement，因此当前必须关闭 CUDA Graph。 +- 推测解码/MTP 当前会显式拒绝，直到每个 proposed token 都能参与 Stage 1 元数据生成和语义 cache 更新。 +- sparse prefill 当前仅对 MiniCPM4.1 生产形状（`QH=32`、`KVH=2`、`D=128`、`block_size=64`）、FP16/BF16 cache、初始 causal batch-1 请求启用；其他布局回退到 dense prefill。 +- 本次 32K 验收使用单个调度器级 32K prefill 请求（内部仍按 4,096 token 的 sparse batch 处理），因此没有覆盖调度器级多 chunk prefill。使用 `--max-num-batched-tokens 8192 --enable-chunked-prefill` 时，现有 FlashAttention 和 InfLLM-V2 服务都会在第三个 8K prefill chunk 触发 Paddle `optional::get()` 未初始化断言；这说明问题位于共享的 chunked-prefill 路径，而不是 sparse Stage 1/Stage 2。本实现没有修改该公共路径，32K 对照统一使用 `--max-num-batched-tokens 32768`。 + +每层、每个 physical block 的理论额外语义-cache 字节数为： + +```text +dtype_bytes * kv_heads * head_dim * ( + block_size / kernel_stride + + block_size / (4 * kernel_stride) +) +``` + +对 MiniCPM4.1 的 BF16、`kv_heads=2`、`head_dim=128` 和默认窗口，该理论值为 2,560 bytes/block/layer。这是由形状推导的静态费用，不是 INF-3 实测峰值显存。 + +## 构建与回归 + +InfLLM-V2 修改了 CUDA custom ops，运行算子测试前需要重新构建。下列示例使用项目 `.venv`；请按实际 CUDA 安装路径和 GPU compute capability 调整 `CUDA_HOME` 与架构列表： + +```shell +export CUDA_HOME=/path/to/cuda +export PATH="${CUDA_HOME}/bin:${PATH}" +export LD_LIBRARY_PATH="${CUDA_HOME}/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" +export PYTHONPATH="${PWD}" + +MAX_JOBS=32 FD_BUILD_RESUME=1 \ + bash build.sh 0 "${PWD}/.venv/bin/python" false "[86]" +``` + +与本特性直接相关的回归命令： + +```shell +CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. .venv/bin/python -m pytest -q \ + tests/model_executor/test_infllmv2_attention_backend.py \ + tests/operators/test_infllmv2_attention_forward.py +``` + +高层测试覆盖 writer/Stage 1/Stage 2 顺序、post-RoPE query 和语义 cache 生命周期；GPU 算子测试覆盖跨 page 语义窗口、每请求/每 KV head 选块、paged mapping、短上下文闭环与非法元数据。该特性不要求文档测试，也不应修改 `tests/model_executor/test_minicpm41.py`。 + +## Dense/Sparse 长上下文对照 + +`benchmarks/benchmark_infllmv2.py` 有七个子命令： + +- `selector`：检查 raw selector trace 并生成未绑定负载的命中率摘要；该摘要仅供查看，不能附加到 sparse 计时结果。 +- `sparse-diagnostic`：对启用 trace 的 sparse 服务执行一次诊断负载，自动向 `benchmark_serving.py` 传入 `--no-warmup`，并在 selector diagnostic JSON 中将新生成的 trace 与精确 prompt-token 负载绑定。 +- `run`：调用 `benchmark_serving.py` 收集吞吐、TTFT 和 TPOT，同时采样指定 GPU 的显存；sparse 计时运行必须通过 `--selector-diagnostic` 提供已绑定的诊断 JSON。 +- `report`：仅在 dense/sparse workload 指纹完全一致且两次运行的 physical GPU UUID 相同时生成对照 JSON 和 Markdown 表格。 +- `operators`：使用 CUDA event 分别测量 compressed-K update、Stage 1、Stage 2 和完整 sparse decode 算子链；默认矩阵为 32K/并发 1 与 128K/并发 4。 +- `prefill`：使用 CUDA event 对比完整的 Paddle dense/sparse prefill 路径。 +- `cuda-impl`：对本地 `infllm_v2` PyTorch extension checkout 执行等价 decode 负载。 + +逐算子矩阵应与服务计时分开运行： + +```shell +CUDA_VISIBLE_DEVICES=7 .venv/bin/python benchmarks/benchmark_infllmv2.py operators \ + --gpu-index 7 \ + --scenario 32768:1 \ + --scenario 131072:4 \ + --output runs/bench/infllmv2/operators.json +``` + +JSON 会为每个场景记录 CUDA mean/median/P10/P90/P99、host wall time、源码哈希、选中 token 比例和持久化 workspace 字节数，并校验 Stage 2 的四个输出均与调用方缓冲区共享存储。该合成 benchmark 用于隔离 kernel 成本，不能替代 TTFT/TPOT 服务对照。 + +sparse prefill 使用独立的完整链路 benchmark： + +```shell +CUDA_VISIBLE_DEVICES=7 PYTHONPATH=. .venv/bin/python \ + benchmarks/benchmark_infllmv2.py prefill \ + --device gpu:0 \ + --context-length 16384 \ + --context-length 32768 \ + --output runs/bench/infllmv2/sparse_prefill.json +``` + +sparse 计时包含 compressed-K 更新、dense prefix、Stage 1、paged-cache gather、两个 FlashAttention 分区及 LSE 合并。2026-08-18 在 NVIDIA RTX A6000、BF16、batch 1 上，32K 实测 dense 90.300 ms、sparse 46.005 ms，即 **1.963x**；16K 为 21.256 ms 对 18.654 ms，即 1.139x，说明该优化面向长上下文。可追溯结果位于 `runs/bench/infllmv2/sparse_prefill_20260818.json`。 + +### 32K/batch 1 理论收益上限 + +默认配置下，sparse Stage 2 选择 96 个、每个 64 token 的 page。在 32K 上实际参与 Stage 2 的是 6,144 / 32,768 个可见 token，即 18.75%。即便假设完全受显存带宽限制且没有任何固定开销，Stage 2 的 token 缩减理论上限也只有 `32768 / 6144 = 5.333x`。端到端 decode 不可能超过该值，并受到更严格的 Amdahl 上限约束： + +```text +speedup <= 1 / ((1 - f_stage2) + f_stage2 / 5.333) +``` + +其中 `f_stage2` 是 dense decode 中 Stage 2 attention 的耗时占比。Stage 1、compressed-K update、split merge、launch 开销、MLP/MoE 以及 batch 1 的 GPU 欠占用均不属于理想 token 缩减，都会压低实测收益。128K 的选中 token 比例更低，但 Stage 1 需要扫描更多语义窗口，因此不能把 token 比例直接当作实测加速比。 + +### 选块命中率 + +命中率是所有“请求/query token/KV head”样本的 dense-reference block recall 微平均： + +```text +sum(|selected_blocks ∩ reference_blocks|) / sum(|reference_blocks|) +``` + +`reference_blocks` 应使用同一 post-RoPE query 和同一 paged K cache 的精确 dense attention 生成，并按 GQA 组聚合到 KV head 粒度。raw sample 格式为： + +```json +{ + "samples": [ + { + "selected_blocks": [0, 12, 37], + "reference_blocks": [0, 9, 37] + } + ] +} +``` + +数组中的 block id 必须非负且不重复，`reference_blocks` 不得为空。该样例只说明 JSON 格式，其数字不是模型实测结果。 + +后端可以在独立诊断运行中直接生成这些样本。请先设置 trace 环境变量，再启动专用的 `INFLLMV2_ATTN` 诊断服务： + +```shell +export FD_INFLLMV2_SELECTOR_TRACE_PATH="${PWD}/runs/bench/infllmv2/selector_samples.json" +export FD_INFLLMV2_SELECTOR_TRACE_RANK=0 +export FD_INFLLMV2_SELECTOR_TRACE_LAYER=0 +export FD_INFLLMV2_SELECTOR_TRACE_MAX_SAMPLES=16 +``` + +启用 trace 时必须关闭 CUDA Graph。后端在 Stage 1 之后使用同一 post-RoPE query、`block_tables` 和已写 paged K cache 计算 FP32 dense softmax；对共享一个 KV head 的 query heads 汇总每个逻辑块的 attention mass，再取与 `selected_count` 相同数量的最高质量块作为 reference。trace 会产生 GPU 同步、device-to-host 复制和 dense attention 额外计算，因此诊断延迟不得作为性能结果。诊断前 trace 路径必须不存在，服务也会显式拒绝覆盖已有文件。 + +```shell +mkdir -p runs/bench/infllmv2 + +.venv/bin/python benchmarks/benchmark_infllmv2.py sparse-diagnostic \ + --base-url http://127.0.0.1:8180 \ + --model MiniCPM4.1-8B \ + --tokenizer "${MODEL_PATH}" \ + --input-len 32768 \ + --output-len 128 \ + --num-prompts 4 \ + --max-concurrency 1 \ + --trace-path runs/bench/infllmv2/selector_samples.json \ + --output runs/bench/infllmv2/selector_diagnostic.json +``` + +`sparse-diagnostic` 会通过传入 `--no-warmup` 自动关闭常规 warm-up 请求，不会生成需要上报的 warm-up 结果。它会把实际生成的 prompt-token 摘要、workload 指纹和 trace 来源一起记录；之后的 sparse 计时 `run` 会拒绝负载不匹配的诊断文件。 + +也可以单独汇总 raw trace 便于查看，但该未绑定摘要不能被 `run` 接受： + +```shell +.venv/bin/python benchmarks/benchmark_infllmv2.py selector \ + --samples runs/bench/infllmv2/selector_samples.json \ + --output runs/bench/infllmv2/selector_summary.json +``` + +### 固定负载运行 + +诊断完成后必须完全停止启用 trace 的 sparse 服务。raw trace 文件需要保留，因为绑定后的诊断会按路径和 SHA256 再次校验它；但任何计时前都必须禁用 trace 并重新启动服务： + +```shell +unset FD_INFLLMV2_SELECTOR_TRACE_PATH +``` + +诊断服务和计时服务必须是不同进程。请确认计时用 `INFLLMV2_ATTN` 服务启动时没有 trace-path 环境变量，并在整个计时期间保持 trace 关闭。 + +正式计时对照中，先使用 `FD_ATTENTION_BACKEND=FLASH_ATTN` 启动 dense 服务，运行： + +```shell +.venv/bin/python benchmarks/benchmark_infllmv2.py run \ + --variant dense \ + --base-url http://127.0.0.1:8180 \ + --model MiniCPM4.1-8B \ + --tokenizer "${MODEL_PATH}" \ + --gpu-index 0 \ + --input-len 32768 \ + --output-len 128 \ + --num-prompts 4 \ + --max-concurrency 1 \ + --output runs/bench/infllmv2/dense.json +``` + +停止 dense 服务，在同一块空闲 GPU 上使用 `FD_ATTENTION_BACKEND=INFLLMV2_ATTN` 和完全相同的服务参数启动未启用 trace 的 sparse 服务，再运行： + +```shell +.venv/bin/python benchmarks/benchmark_infllmv2.py run \ + --variant sparse \ + --base-url http://127.0.0.1:8180 \ + --model MiniCPM4.1-8B \ + --tokenizer "${MODEL_PATH}" \ + --gpu-index 0 \ + --input-len 32768 \ + --output-len 128 \ + --num-prompts 4 \ + --max-concurrency 1 \ + --selector-diagnostic runs/bench/infllmv2/selector_diagnostic.json \ + --output runs/bench/infllmv2/sparse.json + +.venv/bin/python benchmarks/benchmark_infllmv2.py report \ + --dense-result runs/bench/infllmv2/dense.json \ + --sparse-result runs/bench/infllmv2/sparse.json \ + --output runs/bench/infllmv2/report.json +``` + +`sparse-diagnostic`、dense `run` 和 sparse `run` 的负载参数必须完全一致，包括 model、tokenizer、endpoint、seed、请求数、输入/输出长度、request rate 和并发。脚本会校验实际生成的 prompt tokens，并拒绝不匹配的负载指纹。 + +`--gpu-index` 是 `nvidia-smi --id` 使用的 physical GPU 索引。如果设置 `CUDA_VISIBLE_DEVICES=7`，这里应传 `--gpu-index 7`，而不是进程内的逻辑索引 0。测试时应保证该 GPU 没有其他负载，dense/sparse 之间完全重启服务。 + +每个 `run` 输出中的 `gpu_memory.gpu_uuid` 标识 physical device，`baseline_mib` 是请求开始前已启动服务的显存基线，`peak_mib` 是 GPU 总已用显存峰值，`peak_delta_mib` 是两者差值。`report` 会拒绝 GPU UUID 不同的 dense/sparse 输入，并在结果中保留 UUID、baseline、peak 和 peak delta；验收报告必须完整保留这些数据。 + +### 保留的 P0/P2 检查点结果 + +最终 CUDA event 矩阵保存在 `runs/bench/infllmv2/tmp/p3_operator_matrix.json`。它使用 MiniCPM4.1 生产形状（`QH=32`、`KVH=2`、`D=128`、BF16），在 NVIDIA RTX A6000 上预热 50 次、测量 200 次。下表均为设备侧中位数。 + +| 场景 | Compressed-K update (us) | Stage 1 (us) | Stage 2 (us) | 完整 sparse 算子链 (us) | Stage 2 持久化字节数 | +| --- | ---: | ---: | ---: | ---: | ---: | +| 32K / 并发 1 | 49.152 | 191.488 | 111.616 | 380.928 | 1,073,152 | +| 128K / 并发 4 | 49.152 | 1,662.976 | 215.040 | 1,939.968 | 4,292,608 | + +P0/P2 修改前的 Stage 2 基线保存在 `runs/bench/infllmv2/tmp/p0_p2_operator_baseline.json`。同一 GPU、相同形状的结果为： + +| 场景 | P0 前 Stage 2 (us) | GQA-tiled Stage 2 (us) | 加速比 | 耗时下降 | +| --- | ---: | ---: | ---: | ---: | +| 32K / 并发 1 | 135.168 | 111.616 | 1.211x | 17.42% | +| 128K / 并发 4 | 367.616 | 215.040 | 1.710x | 41.50% | + +P0 独立检查点的 Stage 2 实测为 112.144 us 和 215.040 us，相对保留基线分别为 1.205x 和 1.710x；32K 与最终 111.616 us 的差异属于运行波动。 + +矩阵同时验证最终输出和三个 FP32 partial tensor 均与调用方传入的存储 alias。这里的字节数是每个 layer-backend 实例的开销：并发 1 时 Stage 2 持久化 1,073,152 bytes，并发 4 时为 4,292,608 bytes；每个实例只在形状变化时分配，不再每次 layer 调用分配。对这个 32 层 checkpoint，对应总量为 32.75 MiB 和 131.00 MiB。P2 检查点的 32K/1 Stage 2 为 112.640 us，P0 检查点为 112.144 us（+0.44%）；完整算子链从 381.952 us 变为 380.928 us（-0.27%），128K/4 Stage 2 均为 215.040 us。该变化处于运行噪声范围，说明 in-place 持久化没有引入热路径回归。P0 与 P2 分步结果分别保存在同一临时目录的 `p0_gqa_tiled_operator.json` 和 `p2_persistent_workspace_operator.json`。 + +本次 checkpoint 的 `max_position_embeddings` 与 `rope_scaling.original_max_position_embeddings` 都是 65,536；128K 服务请求会违反模型契约。因此验收要求的 128K/并发大于 1 对照明确采用合成逐算子测试，不把它表述为端到端模型结果。 + +### Stage 1 选块优化 + +后续 Stage 1 GQA-tiled coarse LSE 与精确 radix TopK 的分步结果分别保存在 +`runs/bench/infllmv2/tmp/stage1_p1_gqa_tiled_result.json` 和 +`runs/bench/infllmv2/tmp/stage1_p2_radix_topk_result.json`，合并后的 CUDA event +结果为 `stage1_p1_p2_optimized.json`；同目录的 Nsight Systems trace 与 SQLite +导出提供逐 kernel 归因。测试使用同一块 RTX A6000 和 BF16 生产形状，CUDA +event 数据预热 20 次、测量 100 次，表中为中位数。 + +| 场景 | 优化前 Stage 1 (us) | 优化后 Stage 1 (us) | Stage 1 加速 | 优化前算子链 (us) | 优化后算子链 (us) | 算子链加速 | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| 32K / 并发 1 | 191.488 | 82.944 | 2.309x | 381.952 | 246.784 | 1.548x | +| 128K / 并发 4 | 1,659.904 | 820.224 | 2.024x | 1,935.360 | 1,052.672 | 1.839x | + +| 场景 | 优化前 coarse LSE (us) | GQA-tiled coarse LSE (us) | 加速 | 优化前 TopK (us) | Radix TopK (us) | 加速 | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| 32K / 并发 1 | 111.155 | 7.635 | 14.558x | 34.643 | 13.978 | 2.478x | +| 128K / 并发 4 | 479.877 | 22.688 | 21.151x | 509.055 | 22.579 | 22.545x | + +在该检查点的 128K/并发 4 下，block-score 是 Stage 1 的主要 kernel（773.271 us， +约占 Stage 1 的 94%），因此成为后续选块瓶颈。该检查点的隔离 extension 通过 +10 个 SM86 GPU 正确性用例，包括生产形状 BF16 语义参考以及 2,048 候选的精确 +tie/padding 用例；backend 回归通过 39 个用例。该检查点完成全量 package +重编译后，packaged SM86 算子回归 15 个用例全部通过。下一节给出 block-score +瓶颈的最终处理结果。 + +### GQA-tiled block-score 与 Stage 2 最终优化 + +后续实现参考了 OpenBMB 官方仓库中的 split-KV 与分布式 softmax 结构；参考仓库 +下载在 `runs/bench/infllmv2/tmp/infllmv2_cuda_impl`，固定 commit 为 +`93cf2ec28e5a7acebe3f0bb7329b6c73a1be91f6`。FastDeploy 仍保留自己的 paged +cache、持久化 workspace 和 in-place custom-op 契约。 + +最终隔离 CUDA-event 矩阵为 +`runs/bench/infllmv2/tmp/stage1_p3_gqa_blockscore/operator_optimized_halfwarp.json`， +使用同一块 RTX A6000、BF16 生产形状，预热 20 次并测量 100 次。 + +| 场景 | Compressed-K update (us) | Stage 1 (us) | Stage 2 (us) | 完整 sparse 算子链 (us) | +| --- | ---: | ---: | ---: | ---: | +| 32K / 并发 1 | 49.152 | 48.128 | 59.392 | 149.504 | +| 128K / 并发 4 | 49.152 | 131.072 | 124.928 | 317.440 | + +相对保留的 P0/P2 矩阵，32K/1 完整算子链从 380.928 降至 149.504 us +（2.548x，耗时下降 60.75%），128K/4 从 1,939.968 降至 317.440 us +（6.111x，耗时下降 83.64%）。相对紧邻的 coarse-LSE/radix-TopK 检查点，Stage 1 +在 32K/1 与 128K/4 分别加速 1.723x 和 6.258x，Stage 2 分别加速 1.879x 和 +1.787x。 + +同一临时目录保存了 Nsight Systems report、SQLite 导出和生成的 kernel-summary +CSV。下表按 kernel 平均耗时比较上一检查点与最终实现。 + +| Kernel | 32K 优化前/最终 (us) | 加速 | 128K/4 优化前/最终 (us) | 加速 | +| --- | ---: | ---: | ---: | ---: | +| Fine block-score | 58.585 / 13.370 | 4.382x | 773.271 / 82.341 | 9.391x | +| Stage 2 partial | 70.386 / 40.153 | 1.753x | 152.721 / 95.404 | 1.601x | +| Stage 2 merge | 37.548 / 15.059 | 2.493x | 72.025 / 25.997 | 2.770x | + +当前源码的隔离 extension 已通过 14/14 个 GPU 用例，覆盖 FP32/FP16/BF16 paged +Stage 2、生产 BF16 GQA Stage 1/2、精确 2,048-way tie、复用 buffer 尾部和非法 +契约。完整重编译后的生产 package 与配对服务结果在下文继续报告。 + +在模型支持的 32K 服务负载上，保留的 P0/P2 检查点报告 +`runs/bench/infllmv2/tmp/p3_report_32k_current.json` 对比了从当时同一 packaged +build 启动、服务参数完全相同的 dense 与 sparse 服务；它早于上述 Stage 1 选块 +优化。两份输入的 workload ID 均为 +`b0fd4419f876adfa40e4d92a6e46d56c95c0f5f37e0539abd1a32c8208f6c078`， +prompt-token hash 与 physical GPU UUID 相同，且已通过 `report` 命令校验。 + +| 指标 | Dense | P0/P2 检查点 sparse | Sparse / Dense | +| --- | ---: | ---: | ---: | +| Selector block hit rate | — | 64.4531%（990 / 1536） | — | +| Request throughput (req/s) | 0.092799 | 0.084076 | 0.906x | +| Output throughput (token/s) | 11.878301 | 10.761664 | 0.906x | +| Total-token throughput (token/s) | 3052.723423 | 2765.747673 | 0.906x | +| Mean / median / P99 TTFT (ms) | 7368.456 / 7362.350 / 7399.758 | 7404.749 / 7389.901 / 7467.811 | 1.005x / 1.004x / 1.009x | +| Mean / median / P99 TPOT (ms) | 26.742 / 26.744 / 26.794 | 35.252 / 35.408 / 35.514 | 1.318x / 1.324x / 1.325x | +| Baseline / peak / peak-delta GPU memory (MiB) | 19493 / 24421 / 4928 | 19417 / 24345 / 4928 | 0.996x / 0.997x / 1.000x | + +服务负载为 4 个确定性 random-token 请求，input 32768、output 128、并发 1、request rate `inf`、seed 2026。服务使用 TP1、BF16、`block_size=64`、默认 sparse 配置 `32/16/64/8192/1/2048`，关闭 prefix cache 和 CUDA Graph，并使用单个 32K prefill chunk。测试环境为 NVIDIA RTX A6000 48 GiB（UUID `GPU-56bb1092-c218-9cd3-ad85-9a2735522d28`）、CUDA toolkit 12.8、Paddle 3.3.1（编译 CUDA 12.6）和 Python 3.12.13。 + +该保留服务报告不是新 Stage 1 kernel 的端到端测量：其中 sparse 输出吞吐为 +dense 的 0.906x，mean TPOT 为 dense 的 1.318x。当前源码的算子链在 32K/1 +缩短 35.39%，在 128K/4 缩短 45.61%。在该检查点，仍需完成 packaged server +重编译和同负载服务复测；已完成的复测见下节。更早的 +`runs/bench/infllmv2/report.json` 仍作为历史结果保留（输出吞吐为 dense 的 +0.311x、mean TPOT 为 dense 的 7.945x），不得将它当成当前 kernel 的结果。 + +### 历史 Stage 1 优化 packaged 服务检查点 + +该检查点使用当时的源码完整重编译了 `fastdeploy_ops` package。安装库当时的 SHA-256 为 +`cb37a51292ed2f92139a199f748a1e8fb014d6c7774284cfdc6bfe535ed2d5e4`， +packaged GPU 算子回归 15/15 通过。构建、正确性、服务与 benchmark 日志均保存在 +`runs/bench/infllmv2/tmp/retest_stage1_20260813/`。 + +通过一致性校验的报告为 +`runs/bench/infllmv2/tmp/retest_stage1_20260813/report_32k.json`。dense 与 sparse +使用相同 workload ID、上述相同 prompt-token hash、相同 physical GPU UUID，且两种 +服务之间完全重启。正式计时的 sparse 服务未开启 selector trace。 + +| 指标 | Dense | Stage 1 检查点 sparse | Sparse / Dense | +| --- | ---: | ---: | ---: | +| Selector block hit rate | — | 59.2448%（910 / 1536） | — | +| Request throughput (req/s) | 0.092539 | 0.087938 | 0.950x | +| Output throughput (token/s) | 11.845043 | 11.256072 | 0.950x | +| Total-token throughput (token/s) | 3044.176118 | 2892.810553 | 0.950x | +| Mean / median / P99 TTFT (ms) | 7403.585 / 7401.545 / 7445.184 | 7414.495 / 7413.212 / 7430.486 | 1.001x / 1.002x / 0.998x | +| Mean / median / P99 TPOT (ms) | 26.710 / 26.717 / 26.747 | 31.071 / 31.058 / 31.133 | 1.163x / 1.162x / 1.164x | +| Baseline / peak / peak-delta GPU memory (MiB) | 19493 / 24421 / 4928 | 19417 / 24345 / 4928 | 0.996x / 0.997x / 1.000x | + +相对保留的 P0/P2 sparse 检查点，mean sparse TPOT 从 35.252 ms 降到 +31.071 ms（下降 11.86%，加速 1.135x），sparse 输出吞吐从 10.761664 提升到 +11.256072 token/s（+4.59%）。重复测量的 dense mean TPOT 仅变化 -0.12%，因此 +sparse 收益不能由更快的 dense 基线解释。Selector 诊断会跟随实际生成的 decode +轨迹；两次诊断具有相同 prompt 指纹，但不保证生成 token 相同，因此不能把本次 +59.24% 与历史 64.45% 直接当成受控的 selector 质量变化。 + +在这个历史检查点，sparse 虽有改善，但在 32K/batch 1 尚未反超 dense：mean TPOT 仍高 +4.361 ms（dense 的 1.163x），输出吞吐为 dense 的 0.950x。5.333x 的 token +缩减上限只适用于 Stage 2 K/V 计算。Dense FlashAttention 本身已高度融合；sparse +decode 还需支付 compressed-K update、block-score/coarse-LSE/TopK 选块、split +merge、32 层 kernel launch，以及相同的 MLP/MoE 成本，batch 1 又无法充分占满 +GPU。本检查点中这些固定成本仍超过 Stage 2 节省的 K/V 读取，因此“稀疏”本身不 +保证端到端更快。 + +### 最终 GQA-tiled package 验证与服务复测 + +完成 block-score 和 Stage 2 最终优化后，再次完整重编译了 SM86 生产 package。 +仓库内 +`fastdeploy/model_executor/ops/gpu/fastdeploy_ops/fastdeploy_ops_pd_.so` 的 +SHA-256 为 +`368174499c47eb991595616988f1c2a65bd6384f8cf8353d366511784883bb94`。 +backend 回归 39/39 通过，packaged GPU operator 15/15 通过；成功构建与测试日志 +保存在 `runs/bench/infllmv2/tmp/stage1_p3_gqa_blockscore/full_build/`。 + +最终 package 的 CUDA-event 结果为 +`runs/bench/infllmv2/tmp/final_gqa_decode_20260813/operators_packaged.json`。 +该文件记录最终源码 hash，验证 Stage 2 四个输出 alias，并给出以下设备侧中位数： + +| 场景 | Compressed-K update (us) | Stage 1 (us) | Stage 2 (us) | 完整 sparse 算子链 (us) | +| --- | ---: | ---: | ---: | ---: | +| 32K / 并发 1 | 49.152 | 37.888 | 59.392 | 149.504 | +| 128K / 并发 4 | 49.152 | 131.072 | 125.952 | 316.416 | + +相对保留的 P0/P2 算子链，32K/1 和 128K/4 分别加速 2.548x 与 6.131x； +相对 P0 前的 Stage 2 基线，最终 packaged Stage 2 分别加速 2.276x 与 2.919x。 +上文隔离结果的 128K Stage 2/算子链为 124.928/317.440 us，package 复测为 +125.952/316.416 us，差异属于运行波动。 + +最终配对服务结果位于 +`runs/bench/infllmv2/tmp/final_gqa_decode_20260813/`。它继续使用相同 workload +ID、prompt-token hash、physical GPU UUID 和服务参数。sparse 服务独立重启并重复 +测量两次；两个正式计时服务都没有开启 selector trace。 + +| 指标 | Dense | 最终 sparse 第 1 次 | 最终 sparse 重复测量 | +| --- | ---: | ---: | ---: | +| Selector block hit rate | — | 57.5521%（884 / 1536） | 使用同一绑定诊断 | +| Request throughput (req/s) | 0.092431 | 0.073803（0.798x） | 0.074985（0.811x） | +| Output throughput (token/s) | 11.831154 | 9.446807（0.798x） | 9.598114（0.811x） | +| Total-token throughput (token/s) | 3040.606665 | 2427.829371（0.798x） | 2466.715343（0.811x） | +| Mean / median / P99 TTFT (ms) | 7394.920 / 7395.463 / 7436.516 | 11214.213 / 11260.205 / 11902.204 | 10500.386 / 10554.044 / 10999.020 | +| Mean / median / P99 TPOT (ms) | 26.876 / 26.880 / 27.012 | 18.315 / 19.624 / 25.340 | 22.248 / 22.148 / 24.921 | +| Baseline / peak / peak-delta GPU memory (MiB) | 19493 / 24421 / 4928 | 19417 / 24345 / 4928 | 19417 / 24345 / 4928 | + +两次 sparse 的客户端观测 mean TPOT 均低于 dense，分别下降 31.85%（1.467x）和 +17.22%（1.208x）；这是本实现首次在同负载 32K 对照中得到 sparse TPOT 低于 dense +TPOT。但它不是所有端到端指标都加速：sparse TTFT 高 42.0%-51.6%，128-token +短输出的 output throughput 仍只有 dense 的 0.798x-0.811x。 + +最终 sparse 流式输出还存在 burst：两次分别有 45.08% 和 37.20% 的 ITL 小于 +1 ms，而 dense 为 0%。排除 burst 后，benchmark 给出的 clean decode rate 为 +dense 37.21 token/s、sparse 30.04/28.26 token/s。因此这里把较低 TPOT 表述为可 +重复的客户端观测延迟并报告区间，不把单次最佳 1.467x 重新解释成稳定的 kernel +吞吐；上文 CUDA-event 与 Nsight 数据才是 kernel 级证据。 + +这也解释了为什么 sparse 理论上有优势，却不保证每个指标都更快：32K 下 Stage 2 +只读取 dense token 的 18.75%，但 selector、merge、launch、prefill/TTFT 以及非 +attention 模型计算仍然存在。完整原始结果与解释汇总在 +`runs/bench/infllmv2/tmp/final_gqa_decode_20260813/summary.md`。 + +### CUDA 恢复后的基线复测（2026-08-16） + +上面的 checkpoint 章节保留为优化历史。该检查点的当前源码基线结果是 +`runs/bench/infllmv2/service_gate_65k8_20260816/summary.md`；该结果来自重新完整 +构建的 SM86 package，没有沿用历史结论。build 产物与仓库实际加载库逐字节一致， +大小为 1,040,001,400 bytes，SHA-256 为 +`05e5b67e43d0da09bcbb884ae5b78fba23b8a3a99e31be69a3216c7a5d00e2c7`。 +当前 package 的聚焦回归 24/24 通过，覆盖跨 page compressed-K、Stage 1 +score/Top-K、paged Stage 2，以及独立 NumPy/Paddle reference。 + +在模型契约内的 65,344-token、batch 8 BF16 负载上（`QH=32`、`KVH=2`、 +`D=128`、选中 96 pages），CUDA-event 中位数为： + +| 实现 | Stage 1 + Top-K (us) | Stage 2 (us) | Sparse 算子链 (us) | Dense 基线 (us) | +| --- | ---: | ---: | ---: | ---: | +| FastDeploy | 349.184 | 318.464 | 650.240 | 785.408 | +| 本地 `infllmv2_cuda_impl` | 318.464 | 208.896 | 573.440 | 1,563.648 | + +FastDeploy 相对自身 dense 的算子门禁为 1.2079x。相对本地子仓库，FastDeploy 的 +Stage 1 + Top-K、Stage 2 和稀疏链分别慢 1.0965x、1.5245x 和 1.1339x。 +两个实现的 dense reference 不同，因此 dense 列不得跨实现直接比较。page 对齐的 +合成算子有效率为 98,304 / 1,045,504 = 9.40255%。绑定服务 trace 中当前 causal +page 只有部分 token 可见，实际参与率为 97,332 / 1,045,556 = 9.30911%，整页预设 +为 9.40208%，差 -0.09296 个百分点。 + +Nsight 显示 Stage 2 main 是最大稀疏 kernel（286.205 us），fine block-score +次之（249.816 us）。本轮只尝试了一个 64-token shared-page Stage 2 结构：它虽然 +改善 32K/batch 1，却使 128K/batch 4 Stage 2 回退 17.1%；Nsight 将回退定位到 +47.3 KiB shared-memory tile 导致的 occupancy 降低，因此该实验已撤回。当前保留 +16-token tile，并继续使用 GQA K/V 共享、split-KV 和 online softmax。 + +随后分别完整重启 dense 与 sparse 服务，以相同的 65,344 输入、128 输出、8 请求、 +并发 8 负载计时。通过校验的 workload ID 为 +`4068415f92b013c216341f8461e29f42c979aa0aa76070b9677844817d68a1fd`。 +完整服务没有超过 dense：输出吞吐 5.439 对 5.531 token/s（0.983x），mean TTFT +108,132 对 105,124 ms（1.029x），mean TPOT 628.660 对 627.799 ms（1.001x）。 +Selector micro-recall 为 53.84%（827/1,536）。 + +### Paged K/V 向量化权威复测（2026-08-17） + +当前源码的权威证据是 +`runs/bench/infllmv2/service_gate_vector_load_20260817/summary.md`。SM86 package +再次完整重建，build 产物与仓库实际加载库逐字节一致，大小为 1,040,001,400 bytes， +SHA-256 为 +`1370f08391426352349d420b9268a769d8960615d32a717f1bb66a349dca59dc`； +源码 header SHA-256 为 +`883dbc1f5bfbfa8748bb83b3436766f7681d034a26de2e1811661aa738ccd6ae`。 +聚焦回归再次 24/24 通过。 + +本轮只改变生产 Stage 2 完整 tile 的 paged K/V 搬运：global-to-shared 标量循环改为 +对齐的 16-byte `uint4` load/store，因果边界的非完整 tile 仍显式走标量 mask 路径。 +Stage 1、CUB radix Top-K、16-token 数据布局、四 warp tensor-core 并行结构、每 split +两个 pages、online softmax 与 split combine 均未改变。与本地子仓库 64-token +FlashAttention/CUTE tile 的差异仅保留在已有 paged-cache occupancy 证据要求之处， +没有改写为更简单的标量 attention。 + +同一 65,344-token、batch 8 BF16 负载，20 次 warmup 加 100 次 CUDA events 的中位数为： + +| 实现 | Stage 1 + Top-K (us) | Stage 2 (us) | Sparse 算子链 (us) | Dense 基线 (us) | +| --- | ---: | ---: | ---: | ---: | +| FastDeploy | 329.728 | 163.840 | 507.424 | 784.384 | +| 本地 `infllmv2_cuda_impl` | 304.112 | 206.848 | 547.840 | 1,553.408 | + +相对 2026-08-16 FastDeploy 基线，Stage 2 提升 1.9438x，稀疏链降低 21.96%； +FastDeploy 相对自身 dense 的门禁提高到 1.5458x。相对本轮重新运行的本地实现， +FastDeploy Stage 2 快 1.2625x、稀疏链快 1.0796x，Stage 1 + Top-K 仍慢 +1.0842x。两个框架的 dense reference 不同，仍不得跨实现比较 dense 列。 + +Nsight Systems 测得 Stage 2 main 为 145.918 us，低于旧值 286.205 us；fine +block-score 现为最大稀疏 kernel（255.564 us），所以下轮应转向 Stage 1，而不是继续 +调整 Stage 2 block 数。合成算子有效参与率为 98,304 / 1,045,504 = 9.40255%； +当前绑定服务 trace 为 97,332 / 1,045,556 = 9.30911%，整页预设为 9.40208%， +相差 -0.09296 个百分点。 + +随后完整重启 dense 与无 trace sparse 服务，并验证相同 workload ID。完整服务仍未 +超过 dense：sparse 输出吞吐 4.832 对 4.915 token/s（0.983x），mean TTFT +133,586 对 131,160 ms（1.018x），mean TPOT 614.727 对 605.495 ms（1.015x）。 +Selector micro-recall 为 53.39%（820/1,536）。算子收益成立，但不能据此声称该长 +prefill、128-token 输出负载具有通用端到端加速。 + +## W4A16 attention 与服务端到端记录（2026-08-25） + +本节采用统一约定：延迟加速比 = `FlashAttention 延迟 / InfLLM-v2 延迟`，吞吐加速比 = +`InfLLM-v2 吞吐 / FlashAttention 吞吐`；因此大于 1 才表示 InfLLM-v2 更快。attention +算子与完整服务必须分开看：前者只计 attention kernel/算子链，后者还包含 W4A16 +Linear、MLP、归一化、采样、调度与 API 开销。 + +### 当前源码的独立 attention 实测 + +环境为独占 RTX A6000（SM86），MiniCPM4.1 生产形状 `QH=32`、`KVH=2`、 +`D=128`、`block_size=64`、batch 1、BF16 Q/K/V 与 BF16 KV cache。Prefill 使用 20 +次 warmup 加 100 次 CUDA-event 计时，decode 使用 50 次 warmup 加 200 次计时。 + +| 阶段 | 上下文 | FlashAttention 中位延迟 | InfLLM-v2 完整链中位延迟 | 加速比 | +| --- | ---: | ---: | ---: | ---: | +| Prefill attention | 16,384 | 20.824 ms | 18.514 ms | 1.125x | +| Prefill attention | 32,768 | 88.638 ms | 45.810 ms | 1.935x | +| Decode attention | 32,768 | 62.464 us | 181.248 us | 0.345x | + +32K decode 的完整链由 update 16.384 us、Stage 1 95.232 us 和 Stage 2 58.368 us +组成。Stage 2 单独相对 dense 为 1.070x，但完整链只有 0.345x，即实际慢 2.902x； +不得用 Stage 2 单项代替完整 decode attention 加速比。原始记录见 +[prefill JSON](../../../runs/bench/infllmv2/w4a16_sm86_tp1_20260825/attention_prefill_16k_32k_b1.json) +与 [decode JSON](../../../runs/bench/infllmv2/w4a16_sm86_tp1_20260825/attention_decode_32k_b1.json)。 + +### W4A16 服务口径与当前状态 + +W4A16 表示 Linear 权重为 INT4、激活为 BF16；attention Q/K/V 和 KV cache 仍是 +BF16。仅比较 attention backend，其他服务参数固定为 TP1、`max_model_len=65536`、 +`block_size=64`、`max_num_seqs=1`、`max_num_batched_tokens=32768`、关闭 prefix +cache、chunked prefill 与 CUDA Graph。负载为每请求 32,768 输入、128 输出，4 个请求， +并发 1，seed 2026。 + +量化顺序严格为：读取原始 BF16 checkpoint -> 服务启动期执行 `--quantization wint4` +在线 INT4 权重量化 -> worker 就绪 -> 才开始 warmup 和计时。因此请求 E2E 不含一次性 +启动量化时间；W4A16 对被量化权重的理论位宽压缩是 4x，实际模型/显存压缩还需计入 +scale、未量化参数与 BF16 KV cache。 + +FlashAttention 当前基线成功：平均 TTFT 7,802.485 ms、TPOT 14.034 ms、输出吞吐 +13.345 token/s。InfLLM-v2 在计时前的 32K warmup 已激活 sparse prefill，随后在 +LM head 报 `CUBLAS_STATUS_INVALID_VALUE` 并退出；关闭 overlap schedule 后相同。因此 +当前源码没有可报告的 InfLLM-v2 W4A16 服务加速比，不能将失败样本记为 0x，也不能用 +attention 算子比值冒充 E2E。证据见 [FlashAttention 原始结果](../../../runs/bench/infllmv2/w4a16_sm86_tp1_20260825/serving_flash_attn_32k_128_b1.json) +和 [InfLLM-v2 失败记录](../../../runs/bench/infllmv2/w4a16_sm86_tp1_20260825/serving_infllmv2_failure.json)。 + +### 已保存的 W4A16 成功服务对照（2026-08-19，历史 checkpoint） + +为保留可计算的服务数据，下面单独列出仓库内已成功完成的配对记录。两端使用同一个 +预先离线量化完成的 W4A16 checkpoint、同一 RTX A6000、每请求 65,344 输入和 16 +输出、2 请求、并发 1；两份输入 payload SHA-256 都是 +`1100aec67962ce4cf3cf5e9bc18ab7c1bcc44ff06bbf71d80b058322b4a02aa8`。该结果不代表 +2026-08-25 当前源码，因为当前代码已不再支持这个预量化 checkpoint 的加载路径。 + +| 服务指标 | FlashAttention | InfLLM-v2 | InfLLM-v2 加速比 | +| --- | ---: | ---: | ---: | +| Prefill 代理：mean TTFT | 21,628.915 ms | 12,950.952 ms | 1.670x | +| Decode 代理：mean TPOT | 18.143 ms | 35.478 ms | 0.511x | +| 推导 mean E2E/request | 21,901.054 ms | 13,483.121 ms | 1.624x | +| 实测总时长 | 43.820 s | 26.983 s | 1.624x | +| 输出吞吐 | 0.730 token/s | 1.186 token/s | 1.624x | +| clean decode 吞吐 | 53.332 token/s | 28.186 token/s | 0.529x | + +其中 `mean E2E = mean TTFT + (16 - 1) * mean TPOT`。相同 token 总数下，总时长比与 +输出吞吐比一致。结果说明 sparse prefill 带来明显 TTFT 收益，但 decode 约慢 1.956x； +本负载输出很短，prefill 收益仍使整体 E2E 达到约 1.624x。原始记录为 +[FlashAttention](../../../runs/bench/infllmv2/w4a16_64k_b1_gpu7_20260819/flash_attn_64k_single_prefill.json) +和 [InfLLM-v2](../../../runs/bench/infllmv2/w4a16_64k_b1_gpu7_20260819/infllmv2_64k_single_prefill.json)， +统一计算结果见 [summary.json](../../../runs/bench/infllmv2/w4a16_sm86_tp1_20260825/summary.json)。 + +## 验收表 + +| 任务 | 代码/证据 | 文档编写时状态 | +| --- | --- | --- | +| INF-1：按请求和 head 的 Stage 1 | `infllmv2_update_compressed_k`、`infllmv2_select_blocks`；backend/operator tests | 当前 package 聚焦回归 24/24：backend 15/15、packaged GPU 9/9 | +| INF-2：post-RoPE + paged cache Stage 2 | `decoder_write_cache_with_rope` 后调用 `infllmv2_attention_forward`；顺序与 paged mapping 测试 | 通过：writer 顺序、已写 cache、跨 page 和 paged mapping 均已回归 | +| P0：GQA-tiled Stage 2 | 生产形状 CUDA 路径与 NumPy oracle 算子测试 | 通过；对齐 paged K/V 搬运使当前 65K/8 Stage 2 降至 163.840 us，本地实现为 206.848 us | +| P2：Stage 2 持久化缓冲区 | 后端 workspace 复用与四输出 in-place alias 测试 | 通过；Stage 2 热路径不再调用 `empty`/`full`/`zeros` | +| Stage 1-1：GQA-tiled coarse LSE | split partial/merge kernel、持久化 workspace、语义参考 GPU 测试与 Nsight trace | 通过；32K/1 提升 14.558x，128K/4 提升 21.151x | +| Stage 1-2：精确 radix TopK | 复合 key CUB selector 与 2,048 候选 tie/padding GPU 测试 | 通过；32K/1 提升 2.478x，128K/4 提升 22.545x | +| Stage 1-3：GQA-tiled block-score | 8-block K tile 在 16 个 query heads 间复用；生产 GPU oracle 测试与 Nsight trace | 通过；相对上一检查点，32K/1 提升 4.382x、128K/4 提升 9.391x | +| P3：校准后的长上下文对照 | 当前 CUDA-event 门禁、同形状子仓库复测、Nsight 与完整重启的 65K/8 服务报告 | 算子门禁以 1.5458x 通过；当前完整服务输出吞吐仍为 dense 的 0.983x，因此不声称通用服务加速 | +| 实现文档 | 本文、根目录 `rfc.md` 与当前复测 summary | 已提供可追溯的当前源码结果，2026-08-17 章节取代此前性能检查点 | + +## 参考 + +- [InfLLM-V2 论文](https://arxiv.org/abs/2509.24663) +- [OpenBMB InfLLM-V2 CUDA 实现](https://github.com/OpenBMB/infllmv2_cuda_impl) +- [MiniCPM4.1-8B](https://huggingface.co/openbmb/MiniCPM4.1-8B) +- [AngelSlim/Hy3-GGUF](https://huggingface.co/AngelSlim/Hy3-GGUF) +- [llama.cpp PR #25395](https://github.com/ggml-org/llama.cpp/pull/25395) + +Hy3-GGUF 和 llama.cpp PR #25395 是用户指定的低比特/Hy3 工程参考，不是本 InfLLM-V2 选块算法或性能数据的来源。 diff --git a/docs/zh/features/thinking_budget.md b/docs/zh/features/thinking_budget.md index 70e58d0c789..51f51907813 100644 --- a/docs/zh/features/thinking_budget.md +++ b/docs/zh/features/thinking_budget.md @@ -13,15 +13,27 @@ ## 工作原理 -1. **请求侧预计算（DataProcessor）**：当请求中包含 `thinking_budget`，会基于 prompt 的 token ids 计算是否已进入思考段、是否已结束，以及已有的思考长度。 +1. **请求侧预计算（DataProcessor）**：当请求中包含 `thinking_budget`，会基于 prompt 的 token ids 计算是否已进入思考段、是否已结束，以及已有的思考长度。单token标记可由 `DataProcessor` 预计算；多 token 标记由logits processor回放，以保留未完成的标记前缀。 2. **每步更新**：解码过程中跟踪 `last_token_id` 与 `tokens_after_start`。 3. **预算约束**：达到预算后，默认直接强制 ``；如果配置了 `think_stop_sentence`，则先逐 token 强制输出该文案，再输出 ``。 ## 前置要求 -- 模型需提供有效的 `think_start_id`、`think_end_id`（来自 `ModelConfig`）。 -- 若其中任意 id 无效，处理器会禁用，`thinking_budget` 不生效。 +- 模型必须提供有效的 `think_start_id`、`think_end_id`，或完整的 `think_token_sequences` 配置。 +- 两种模式都未配置时，通用处理器禁用，`thinking_budget` 不生效；显式声明多token marker能力的模型若推导失败，则启动直接失败。 + +### 多 token 思考标记（`think_token_sequences`） + +部分模型（如 MiniCPM4.1）的 tokenizer 词表中没有单 token 的 ``/``， +思考标记会被切成多个 token，且序列随上下文不同（词首与文中变体不同）。此类模型显式提供 +tokenizer 序列构建 hook；引擎在启动时调用该 hook，并通过 +`ModelConfig.think_token_sequences`（结构：`{"start": [[...]], "end": [[...]], "forced_end": [...]}`） +传递给处理器： + +- 思考段检测按 token 序列匹配，序列前缀不计入预算； +- 预算达到后逐 token 强制输出必填的 `forced_end` 序列，并在下一 decode step 校验实际 token； +- 是否推导由显式模型 hook 决定，而不是宽泛的 `REASONING` 类别；构建或配置错误会直接阻止启动。 ## 请求参数 diff --git a/docs/zh/get_started/minicpm41.md b/docs/zh/get_started/minicpm41.md new file mode 100644 index 00000000000..1cb0db4ac23 --- /dev/null +++ b/docs/zh/get_started/minicpm41.md @@ -0,0 +1,248 @@ +[English](../../get_started/minicpm41.md) + +# MiniCPM4.1-8B模型 + +本文档讲解如何使用FastDeploy部署MiniCPM4.1-8B BF16模型或启用在线 WINT4/WINT8 量化。在开始部署前，请确保硬件环境满足如下条件： + +- GPU驱动 >= 535 +- CUDA >= 12.3 +- CUDNN >= 9.5 +- Linux X86_64 +- Python >= 3.10 +- 48GB NVIDIA GPU 1卡 + +安装FastDeploy方式参考[安装文档](./installation/README.md)。 + +## 最小验收准备 + +进入 FastDeploy 仓库，只需指定本地模型目录和测试 GPU；脚本会自动使用项目 `.venv` 并配置其 CUDA 运行库： + +```shell +cd /path/to/FastDeploy +export MODEL_PATH=/path/to/MiniCPM4.1-8B +export CUDA_VISIBLE_DEVICES=0 +``` + +## 准备模型 + +### 1. 手动下载 （可选） + +MiniCPM4.1-8B使用Hugging Face Torch safetensors格式。执行如下命令下载模型： + +```shell +export MODEL_PATH=/path/to/MiniCPM4.1-8B +hf download openbmb/MiniCPM4.1-8B --local-dir "${MODEL_PATH}" +``` + +模型配置中的 `pad_token_id` 为2。部署前需要确保 `tokenizer_config.json` 中的 `pad_token` 为token ID 2对应的 ``，可执行如下命令完成配置： + +```shell +MODEL_PATH="${MODEL_PATH}" python - <<'PY' +import json +import os +from pathlib import Path + +config_path = Path(os.environ["MODEL_PATH"]) / "tokenizer_config.json" +config = json.loads(config_path.read_text(encoding="utf-8")) +config["pad_token"] = "" +config_path.write_text( + json.dumps(config, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", +) +PY +``` + +### 2. 自动下载 + +将模型名称设置为 `openbmb/MiniCPM4.1-8B`： + +```shell +export MODEL_PATH=openbmb/MiniCPM4.1-8B +``` + +>💡 **注意**：如果当前目录下不存在 `--model` 指定的路径，FastDeploy会根据模型名称查询AIStudio预置模型；查询成功后自动下载，并支持断点续传。默认下载源为AIStudio，可通过 `FD_MODEL_SOURCE` 和 `FD_MODEL_CACHE` 配置下载源及缓存目录，详情参阅[模型下载](../../supported_models.md)。如果已经下载模型，也可以将 `MODEL_PATH` 设置为本地模型目录。 + +## 编译与最小回归 + +### 编译 CUDA 算子 + +完成 CUDA 算子编译和 InfLLM-V2 符号检查： + +```shell +bash tests/benchmarks/test_minicpm41.sh build +``` + +脚本内部调用 `build.sh`。其他 GPU 可通过 `FD_BUILDING_ARCS` 调整架构。验收成功时输出 `PASS build`，并确认产物 `fastdeploy/model_executor/ops/gpu/fastdeploy_ops/fastdeploy_ops_pd_.so` 包含以下符号： + +- `infllmv2_update_compressed_k` +- `infllmv2_select_blocks` +- `infllmv2_attention_forward` + +### MiniCPM4.1、thinking 与量化测试 + +测试模型注册、权重映射、多 token thinking、混合思考模式、WINT4/WINT8 在线量化和 InfLLM-V2 后端： + +```shell +bash tests/benchmarks/test_minicpm41.sh unit +``` + +运行编译后的 CUDA 算子正确性测试： + +```shell +bash tests/benchmarks/test_minicpm41.sh operators +``` + +通过标准为命令退出码0。 + +### E2E 服务测试 + +一条命令依次验收 BF16、WINT4、WINT8 和 InfLLM-V2： + +```shell +bash tests/benchmarks/test_minicpm41.sh e2e +``` + +E2E 会自动选择临时端口、启动和清理服务，并在每种模式运行5个 `test_minicpm41_serving.py` 用例。每种模式显示 `5 passed` 即通过。 + +如需从编译到 E2E 全部一次执行： + +```shell +bash tests/benchmarks/test_minicpm41.sh +``` + +## 启动服务 + +>💡 **注意**：以下命令使用单卡BF16模型，并关闭prefix caching。 + +执行如下命令启动服务，其中启动命令配置方式参考[参数说明](../parameters.md)。 + +```shell +export CUDA_VISIBLE_DEVICES=0 +export FD_ATTENTION_BACKEND=FLASH_ATTN + +python -m fastdeploy.entrypoints.openai.api_server \ + --model "${MODEL_PATH}" \ + --served-model-name MiniCPM4.1-8B \ + --port 8180 --engine-worker-queue-port 8182 \ + --cache-queue-port 8183 --metrics-port 8181 \ + --tensor-parallel-size 1 \ + --max-model-len 8192 \ + --max-num-seqs 1 \ + --max-num-batched-tokens 128 \ + --no-enable-prefix-caching +``` + +### 在线 WINT4/WINT8 + +继续使用同一个原始 BF16 checkpoint，并增加 `--quantization wint4` 或 `--quantization wint8`。FastDeploy 会在 BF16 权重加载后于内存中逐个量化 Linear 权重；MiniCPM4.1 的这条路径不需要、也不支持预先转换好的量化 checkpoint。其余服务参数与上面的 BF16 命令相同。 + +完成前面的最小验收准备后，启动 WINT4： + +```shell +bash scripts/run_minicpm41_wint_server.sh wint4 "${MODEL_PATH}" +``` + +将 `wint4` 替换为 `wint8` 即可启用在线 INT8 权重量化。该脚本要求 `MODEL_PATH` 指向本地模型目录，并自动使用项目 `.venv`、CUDA 运行库及单卡安全默认参数。额外的 Server 参数可继续追加在命令末尾。若端口冲突，应同时为 API、metrics、engine queue 和 cache queue 更换一组未占用端口。公共 WINT 行为见[在线量化](../quantization/online_quantization.md)。 + +## 用户发起服务请求 + +执行启动服务指令后，当终端打印如下信息，说明服务已经启动成功。 + +```shell +INFO api_server.py[line:1030] Launching metrics service at http://0.0.0.0:8181/metrics +INFO api_server.py[line:1033] Launching chat completion service at http://0.0.0.0:8180/v1/chat/completions +INFO api_server.py[line:1034] Launching completion service at http://0.0.0.0:8180/v1/completions +[INFO] Starting gunicorn 26.0.0 +[INFO] Listening at: http://0.0.0.0:8180 +[INFO] Application startup complete. +``` + +FastDeploy提供服务探活接口，用以判断服务的启动状态，执行如下命令返回 `HTTP/1.1 200 OK` 即表示服务启动成功。 + +```shell +curl -i http://0.0.0.0:8180/health +``` + +通过如下命令进行服务请求。`enable_thinking=false` 表示关闭思考模式。 + +```shell +curl -X POST "http://0.0.0.0:8180/v1/chat/completions" \ +-H "Content-Type: application/json" \ +-d '{ + "model": "MiniCPM4.1-8B", + "messages": [ + {"role": "user", "content": "把李白的静夜思改写为现代诗"} + ], + "temperature": 0, + "top_p": 1, + "max_tokens": 64, + "stream": false, + "chat_template_kwargs": {"enable_thinking": false} +}' | jq --indent 4 . +``` + +response如下： + +```json + +{ + "object": "chat.completion", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "好的，这是将李白的《静夜思》改写为现代诗的版本：\n\n**《夜思》**\n\n窗外，月光\n如流水般倾泻\n\n白茫茫\n铺满了冰冷的窗棂\n\n远方的家\n在记忆深处闪烁\n\n小小的床\n" + }, + "logprobs": null, + "draft_logprobs": null, + "prompt_logprobs": null, + "finish_reason": "length", + "speculate_metrics": null + } + ], + "usage": { + "prompt_tokens": 29, + "total_tokens": 93, + "completion_tokens": 64, + "prompt_tokens_details": { + "cached_tokens": 0, + "image_tokens": 0, + "video_tokens": 0 + }, + "completion_tokens_details": { + "reasoning_tokens": 0, + "image_tokens": 0 + } + } +} +``` + +FastDeploy服务接口兼容OpenAI协议，可以通过如下Python代码发起服务请求。 + +```python +import openai + +host = "0.0.0.0" +port = "8180" +client = openai.Client(base_url=f"http://{host}:{port}/v1", api_key="null") + +response = client.chat.completions.create( + model="MiniCPM4.1-8B", + messages=[ + {"role": "system", "content": "I'm a helpful AI assistant."}, + {"role": "user", "content": "把李白的静夜思改写为现代诗"}, + ], + temperature=0, + top_p=1, + max_tokens=64, + stream=True, + extra_body={"chat_template_kwargs": {"enable_thinking": False}}, +) +for chunk in response: + if chunk.choices[0].delta.content: + print(chunk.choices[0].delta.content, end="") +print("\n") +``` + +MiniCPM4.1支持思考与非思考两种模式。将请求中的 `chat_template_kwargs.enable_thinking` 设置为 `true` 即可开启思考模式，还可以通过 `reasoning_max_tokens` 限制本次请求使用的思考token数。 diff --git a/docs/zh/supported_models.md b/docs/zh/supported_models.md index 1424d2320fb..941e51ff112 100644 --- a/docs/zh/supported_models.md +++ b/docs/zh/supported_models.md @@ -38,6 +38,7 @@ python -m fastdeploy.entrypoints.openai.api_server \ |⭐DEEPSEEK|BF16/WINT4|unsloth/DeepSeek-V3.1-BF16;
unsloth/DeepSeek-V3-0324-BF16;
unsloth/DeepSeek-R1-BF16, etc.| |⭐GPT-OSS|BF16/WINT8|unsloth/gpt-oss-20b-BF16, etc.| |⭐GLM-4.5/4.6|BF16/wfp8afp8|zai-org/GLM-4.5-Air;
zai-org/GLM-4.6
 [最佳实践](./best_practices/GLM-4-MoE-Text.md) etc.| +|⭐MiniCPM4.1|BF16/WINT4/WINT8|openbmb/MiniCPM4.1-8B
 [测试教程](./get_started/minicpm41.md)| ## 多模态语言模型列表 diff --git a/fastdeploy/config.py b/fastdeploy/config.py index e2f15cfd014..63c01e4fe3c 100644 --- a/fastdeploy/config.py +++ b/fastdeploy/config.py @@ -111,6 +111,40 @@ def try_match_architecture_defaults( return None +def validate_thinking_token_sequences(value, vocab_size: int) -> Optional[dict]: + """Validate the internal multi-token thinking marker transport contract.""" + if value is None: + return None + if not isinstance(value, dict): + raise ValueError("think_token_sequences must be a dict or null") + expected_keys = {"start", "end", "forced_end"} + if set(value) != expected_keys: + raise ValueError("think_token_sequences must contain exactly start, end, and forced_end") + + for name in ("start", "end"): + sequences = value[name] + if not isinstance(sequences, (list, tuple)) or not sequences: + raise ValueError(f"think_token_sequences.{name} must be a non-empty list of token id sequences") + for sequence in sequences: + if not isinstance(sequence, (list, tuple)) or not sequence: + raise ValueError(f"think_token_sequences.{name} entries must be non-empty token id sequences") + if any(isinstance(token_id, bool) or not isinstance(token_id, int) for token_id in sequence): + raise ValueError(f"think_token_sequences.{name} token ids must be integers") + if any(token_id < 0 or token_id >= vocab_size for token_id in sequence): + raise ValueError( + f"think_token_sequences.{name} token ids must be in vocabulary range [0, {vocab_size})" + ) + + forced_end = value["forced_end"] + if not isinstance(forced_end, (list, tuple)) or not forced_end: + raise ValueError("think_token_sequences.forced_end must be a non-empty token id sequence") + if any(isinstance(token_id, bool) or not isinstance(token_id, int) for token_id in forced_end): + raise ValueError("think_token_sequences.forced_end token ids must be integers") + if any(token_id < 0 or token_id >= vocab_size for token_id in forced_end): + raise ValueError(f"think_token_sequences.forced_end token ids must be in vocabulary range [0, {vocab_size})") + return value + + class MoEPhase: """ The generation phase of the moe. @@ -209,6 +243,7 @@ def __init__( self.is_moe_quantized = False self.max_model_len = 0 self.dtype = "bfloat16" + self._model_info = None self.enable_logprob = False self.max_logprobs = 20 self.logprobs_mode = "raw_logprobs" @@ -286,6 +321,11 @@ def __init__( self.ori_vocab_size = args.get("ori_vocab_size", self.vocab_size) self.think_start_id = args.get("think_start_id", -1) self.think_end_id = args.get("think_end_id", -1) + # Optional multi-token marker mode. The processor validates the exact + # {"start": [[...]], "end": [[...]], "forced_end": [...]} contract. + self.think_token_sequences = validate_thinking_token_sequences( + args.get("think_token_sequences"), self.vocab_size + ) self.im_patch_id = args.get("image_patch_id", -1) self.line_break_id = args.get("line_break_id", -1) self.think_truncate_prompt_ids = args.get("think_truncate_prompt_ids", [-1]) diff --git a/fastdeploy/engine/common_engine.py b/fastdeploy/engine/common_engine.py index 322e880475a..e0ed86dcd75 100644 --- a/fastdeploy/engine/common_engine.py +++ b/fastdeploy/engine/common_engine.py @@ -41,7 +41,7 @@ import fastdeploy.metrics.trace as tracing from fastdeploy.cache_manager.cache_data import CacheStatus -from fastdeploy.config import FDConfig +from fastdeploy.config import FDConfig, validate_thinking_token_sequences from fastdeploy.engine.register_manager import RegisterManager from fastdeploy.engine.request import ( ControlRequest, @@ -124,6 +124,23 @@ def _format_worker_launch_failure_message(paddle_log_dir: str) -> str: return message +def _build_model_thinking_token_sequences(model_config, tokenizer) -> Optional[dict]: + """Invoke a model-owned tokenizer marker builder when the model provides one. + + model_config._model_info is set only after ModelConfig._post_init completed + registry inspection (which already imported the model class), so resolving + the class here never re-imports model modules on the engine side. + """ + if model_config._model_info is None: + return None + model_cls, _ = model_config.registry.resolve_model_cls(model_config.architectures) + builder = getattr(model_cls, "build_thinking_token_sequences", None) + if builder is None: + return None + sequences = builder(tokenizer) + return validate_thinking_token_sequences(sequences, model_config.vocab_size) + + class EngineService: """ Base class containing common engine functionality @@ -2553,12 +2570,20 @@ def _start_worker_service(self): if think_start_id >= 0: self.llm_logger.info(f"Get think_start_id {think_start_id} from vocab.") else: - self.llm_logger.info("No token found in vocabulary, the model can not do reasoning.") + self.llm_logger.info("No single-token marker found in vocabulary.") think_end_id = self.data_processor.tokenizer.get_vocab().get("", -1) if think_end_id >= 0: self.llm_logger.info(f"Get think_end_id {think_end_id} from vocab.") else: - self.llm_logger.info("No token found in vocabulary, the model can not do reasoning.") + self.llm_logger.info("No single-token marker found in vocabulary.") + # MiniCPM4.1: / are multi-token; derive via model hook when single-token ids are absent. + think_token_sequences = None + if think_start_id < 0 or think_end_id < 0: + think_token_sequences = _build_model_thinking_token_sequences( + self.cfg.model_config, self.data_processor.tokenizer + ) + if think_token_sequences is not None: + self.llm_logger.info(f"Get think_token_sequences {think_token_sequences} from tokenizer.") image_patch_id = self.data_processor.tokenizer.get_vocab().get("<|IMAGE_PLACEHOLDER|>", -1) line_break_id = self.data_processor.tokenizer.get_vocab().get("\n", -1) if line_break_id < 0: @@ -2607,6 +2632,7 @@ def _start_worker_service(self): f" --ori_vocab_size {ori_vocab_size}" f" --think_start_id {think_start_id}" f" --think_end_id {think_end_id}" + f" --think_token_sequences '{json.dumps(think_token_sequences)}'" f" --image_patch_id {image_patch_id}" f" --line_break_id {line_break_id}" f" --speculative_config '{self.cfg.speculative_config.to_json_string()}'" diff --git a/fastdeploy/engine/engine.py b/fastdeploy/engine/engine.py index 8614d29999b..d5f4f69f688 100644 --- a/fastdeploy/engine/engine.py +++ b/fastdeploy/engine/engine.py @@ -39,6 +39,7 @@ from fastdeploy.engine.args_utils import EngineArgs from fastdeploy.engine.common_engine import ( EngineService, + _build_model_thinking_token_sequences, _format_worker_launch_failure_message, ) from fastdeploy.engine.expert_service import start_data_parallel_service @@ -570,12 +571,19 @@ def _start_worker_service(self): if think_start_id >= 0: llm_logger.info(f"Get think_start_id {think_start_id} from vocab.") else: - llm_logger.info("No token found in vocabulary, the model can not do reasoning.") + llm_logger.info("No single-token marker found in vocabulary.") think_end_id = self.data_processor.tokenizer.get_vocab().get("", -1) if think_end_id >= 0: llm_logger.info(f"Get think_end_id {think_end_id} from vocab.") else: - llm_logger.info("No token found in vocabulary, the model can not do reasoning.") + llm_logger.info("No single-token marker found in vocabulary.") + think_token_sequences = None + if think_start_id < 0 or think_end_id < 0: + think_token_sequences = _build_model_thinking_token_sequences( + self.cfg.model_config, self.data_processor.tokenizer + ) + if think_token_sequences is not None: + llm_logger.info(f"Get think_token_sequences {think_token_sequences} from tokenizer.") image_patch_id = self.data_processor.tokenizer.get_vocab().get("<|IMAGE_PLACEHOLDER|>", -1) line_break_id = self.data_processor.tokenizer.get_vocab().get("\n", -1) if line_break_id < 0: @@ -645,6 +653,7 @@ def _start_worker_service(self): f" --ori_vocab_size {ori_vocab_size}" f" --think_start_id {think_start_id}" f" --think_end_id {think_end_id}" + f" --think_token_sequences '{json.dumps(think_token_sequences)}'" f" --image_patch_id {image_patch_id}" f" --line_break_id {line_break_id}" f" --think_truncate_prompt_ids '{json.dumps(think_truncate_prompt_ids)}'" diff --git a/fastdeploy/model_executor/layers/attention/__init__.py b/fastdeploy/model_executor/layers/attention/__init__.py index dee4de4df76..8a5331d0642 100644 --- a/fastdeploy/model_executor/layers/attention/__init__.py +++ b/fastdeploy/model_executor/layers/attention/__init__.py @@ -21,6 +21,7 @@ from .dsa_attention_backend import DSAAttentionBackend from .flash_attn_backend import FlashAttentionBackend from .flash_mask_attn_backend import FlashMaskAttentionBackend +from .infllmv2_attention_backend import InfLLMV2AttentionBackend from .mla_attention_backend import MLAAttentionBackend from .moba_attention_backend import PlasAttentionBackend from .native_paddle_backend import PaddleNativeAttnBackend @@ -38,4 +39,5 @@ "PlasAttentionBackend", "FlashMaskAttentionBackend", "DecodeUnifiedAttentionBackend", + "InfLLMV2AttentionBackend", ] diff --git a/fastdeploy/model_executor/layers/attention/infllmv2_attention_backend.py b/fastdeploy/model_executor/layers/attention/infllmv2_attention_backend.py new file mode 100644 index 00000000000..de14480ddc2 --- /dev/null +++ b/fastdeploy/model_executor/layers/attention/infllmv2_attention_backend.py @@ -0,0 +1,1088 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +import math +import os +import weakref +from dataclasses import dataclass, fields +from pathlib import Path +from typing import TYPE_CHECKING, Optional + +import paddle +from paddleformers.utils.log import logger + +from fastdeploy import envs +from fastdeploy.config import FDConfig +from fastdeploy.model_executor.layers.attention.flash_attn_backend import ( + FlashAttentionBackend, + FlashAttentionMetadata, + flash_attn_func, +) +from fastdeploy.model_executor.layers.attention.ops import ( + decoder_write_cache_with_rope, + get_block_shape_and_split_kv_block, + gqa_rope_write_cache, + init_signal_layerwise, + pre_cache_len_concat, +) +from fastdeploy.platforms import current_platform + +if TYPE_CHECKING: + from fastdeploy.model_executor.forward_meta import ForwardMeta + + +_STAGE1_COARSE_WINDOWS_PER_SPLIT = 16 +_STAGE1_MAX_CANDIDATE_BLOCKS = 2048 +_STAGE2_BLOCKS_PER_SPLIT = 2 +_PREFILL_QUERY_CHUNK_SIZE = 4096 +_PREFILL_QUERY_TILE_SIZE = 128 +_SELECTOR_REFERENCE_DEFINITION = ( + "top blocks by exact dense softmax attention mass, summed over the query heads sharing one KV head" +) + + +@dataclass +class InfLLMV2AttentionMetadata(FlashAttentionMetadata): + """Transient InfLLM-V2 metadata for one model forward.""" + + topk_indices: Optional[paddle.Tensor] = None + block_scores: Optional[paddle.Tensor] = None + selected_counts: Optional[paddle.Tensor] = None + compressed_k: Optional[paddle.Tensor] = None + compressed_k2: Optional[paddle.Tensor] = None + kernel_size: int = 32 + kernel_stride: int = 16 + topk: int = 64 + dense_len: int = 8192 + block_size: int = 64 + init_blocks: int = 1 + local_blocks: int = 32 + selected_capacity: int = 128 + + +def _config_value(fd_config: FDConfig, name: str, default): + model_config = fd_config.model_config + sparse_config = getattr(model_config, "sparse_config", None) + if isinstance(sparse_config, dict) and name in sparse_config: + return sparse_config[name] + if hasattr(model_config, name): + return getattr(model_config, name) + if hasattr(fd_config, name): + return getattr(fd_config, name) + return default + + +class InfLLMV2AttentionBackend(FlashAttentionBackend): + """MiniCPM4.1 CUDA backend implementing InfLLM-V2.""" + + __infer_dynamic_dims_fields__ = ["attention_metadata"] + attention_metadata: InfLLMV2AttentionMetadata + + def __init__( + self, + fd_config: FDConfig, + kv_num_heads: int, + num_heads: int, + head_dim: int, + encoder_block_shape_q: int = -1, + decoder_block_shape_q: int = -1, + ) -> None: + super().__init__( + fd_config=fd_config, + kv_num_heads=kv_num_heads, + num_heads=num_heads, + head_dim=head_dim, + encoder_block_shape_q=encoder_block_shape_q, + decoder_block_shape_q=decoder_block_shape_q, + ) + if current_platform.is_cuda() and paddle.device.cuda.get_device_capability()[0] < 8: + raise ValueError("InfLLM-V2 requires NVIDIA compute capability 8.0 or newer.") + + self.kernel_size = int(_config_value(fd_config, "kernel_size", 32)) + self.kernel_stride = int(_config_value(fd_config, "kernel_stride", 16)) + self.topk = int(_config_value(fd_config, "topk", 64)) + self.dense_len = int(_config_value(fd_config, "dense_len", 8192)) + self.init_blocks = int(_config_value(fd_config, "init_blocks", 1)) + self.window_size = int(_config_value(fd_config, "window_size", 2048)) + self.sparse_prefill = bool(_config_value(fd_config, "sparse_prefill", True)) + self.prefill_query_chunk_size = int( + _config_value(fd_config, "prefill_query_chunk_size", _PREFILL_QUERY_CHUNK_SIZE) + ) + + if self.kv_num_heads <= 0 or self.num_heads % self.kv_num_heads != 0: + raise ValueError("InfLLM-V2 num_heads must be divisible by kv_num_heads.") + if self.kernel_size <= 0 or self.kernel_stride <= 0: + raise ValueError("InfLLM-V2 kernel_size and kernel_stride must be positive.") + if self.block_size % self.kernel_stride != 0: + raise ValueError("InfLLM-V2 block_size must be divisible by kernel_stride.") + if self.block_size % (4 * self.kernel_stride) != 0: + raise ValueError("InfLLM-V2 block_size must be divisible by 4 * kernel_stride.") + if self.topk <= 0: + raise ValueError("InfLLM-V2 topk must be positive.") + if self.dense_len < 4 * self.kernel_size: + raise ValueError("InfLLM-V2 dense_len must cover at least one coarse semantic window.") + if self.init_blocks < 0 or self.init_blocks >= self.topk: + raise ValueError("InfLLM-V2 init_blocks must be non-negative and smaller than topk.") + if self.window_size < 0 or self.window_size % self.block_size != 0: + raise ValueError("InfLLM-V2 window_size must be a non-negative multiple of block_size.") + if self.prefill_query_chunk_size <= 0 or self.prefill_query_chunk_size % _PREFILL_QUERY_TILE_SIZE != 0: + raise ValueError("InfLLM-V2 prefill_query_chunk_size must be a positive multiple of 128.") + + max_candidate_blocks = (self.max_seq_len + self.block_size - 1) // self.block_size + if max_candidate_blocks > _STAGE1_MAX_CANDIDATE_BLOCKS: + raise ValueError("InfLLM-V2 supports at most 2048 candidate blocks per request.") + self.local_blocks = self.window_size // self.block_size + self.selected_capacity = max( + self.topk + self.local_blocks, + (self.dense_len + self.block_size - 1) // self.block_size, + ) + if self.selected_capacity * self.block_size > 8192: + raise ValueError("InfLLM-V2 selected capacity may cover at most 8192 tokens.") + + tensor_parallel_size = int(fd_config.parallel_config.tensor_parallel_size) + global_kv_heads = getattr(fd_config.model_config, "num_key_value_heads_list", None) + if global_kv_heads is None: + global_kv_heads = [ + int( + getattr( + fd_config.model_config, + "num_key_value_heads", + self.kv_num_heads * tensor_parallel_size, + ) + ) + ] + if tensor_parallel_size > min(int(value) for value in global_kv_heads): + raise ValueError("InfLLM-V2 does not support tensor parallel KV-head replication.") + if self.speculative_method is not None: + raise ValueError("InfLLM-V2 does not support speculative decoding.") + + graph_opt_config = getattr(fd_config, "graph_opt_config", None) + if graph_opt_config is not None and getattr(graph_opt_config, "use_cudagraph", False): + raise ValueError("InfLLM-V2 requires CUDA Graph to be disabled.") + cache_config = fd_config.cache_config + if getattr(cache_config, "enable_prefix_caching", False): + raise ValueError("InfLLM-V2 does not support prefix caching.") + if getattr(cache_config, "num_cpu_blocks", 0) or getattr(cache_config, "kvcache_storage_backend", None): + raise ValueError("InfLLM-V2 does not support KV-cache offload.") + scheduler_config = getattr(fd_config, "scheduler_config", None) + if self.pd_disaggregation_mode not in (None, "None") or ( + scheduler_config is not None and getattr(scheduler_config, "splitwise_role", "mixed") != "mixed" + ): + raise ValueError("InfLLM-V2 does not support P/D disaggregation.") + if envs.ENABLE_V1_KVCACHE_MANAGER: + raise ValueError("InfLLM-V2 does not support the V1 KV-cache manager.") + + self._compressed_k: Optional[paddle.Tensor] = None + self._compressed_k2: Optional[paddle.Tensor] = None + self._compressed_cache_owner: Optional[weakref.ReferenceType[paddle.Tensor]] = None + self._workspace_key = None + self._topk_indices_ws: Optional[paddle.Tensor] = None + self._block_scores_ws: Optional[paddle.Tensor] = None + self._selected_counts_ws: Optional[paddle.Tensor] = None + self._coarse_lse_ws: Optional[paddle.Tensor] = None + self._coarse_partial_max_ws: Optional[paddle.Tensor] = None + self._coarse_partial_sum_ws: Optional[paddle.Tensor] = None + self._attention_out_ws: Optional[paddle.Tensor] = None + self._partial_acc_ws: Optional[paddle.Tensor] = None + self._partial_max_ws: Optional[paddle.Tensor] = None + self._partial_sum_ws: Optional[paddle.Tensor] = None + self._logged_sparse_activation = False + self._logged_sparse_prefill_activation = False + + trace_path = os.getenv("FD_INFLLMV2_SELECTOR_TRACE_PATH") + self._selector_trace_path = Path(trace_path).resolve() if trace_path else None + self._selector_trace_rank = int(os.getenv("FD_INFLLMV2_SELECTOR_TRACE_RANK", "0")) + self._selector_trace_layer = int(os.getenv("FD_INFLLMV2_SELECTOR_TRACE_LAYER", "0")) + self._selector_trace_max_samples = int(os.getenv("FD_INFLLMV2_SELECTOR_TRACE_MAX_SAMPLES", "16")) + self._selector_trace_samples = [] + self._selector_trace_last_position = {} + self._selector_trace_request_index = {} + self._selector_trace_next_request_index = 0 + if self._selector_trace_path is not None: + if self._selector_trace_rank < 0 or self._selector_trace_layer < 0: + raise ValueError("InfLLM-V2 selector trace rank and layer must be non-negative.") + if self._selector_trace_max_samples <= 0: + raise ValueError("InfLLM-V2 selector trace max samples must be positive.") + if self._selector_trace_path.exists(): + raise FileExistsError(f"InfLLM-V2 selector trace already exists: {self._selector_trace_path}") + if not self._selector_trace_path.parent.is_dir(): + raise FileNotFoundError( + f"InfLLM-V2 selector trace parent directory does not exist: {self._selector_trace_path.parent}" + ) + + def _split_qkv(self, qkv: paddle.Tensor): + if qkv is None: + raise ValueError("InfLLM-V2 sparse decode requires fused qkv input.") + if len(qkv.shape) != 2: + raise ValueError(f"fused qkv must have shape [tokens, width], got {list(qkv.shape)}.") + q_width = self.num_heads * self.head_dim + kv_width = self.kv_num_heads * self.head_dim + expected_width = q_width + 2 * kv_width + if qkv.shape[-1] != expected_width: + raise ValueError(f"fused qkv last dimension must be {expected_width}, got {qkv.shape[-1]}.") + q, k, v = paddle.split(qkv, [q_width, kv_width, kv_width], axis=-1) + return ( + paddle.reshape(q, [-1, self.num_heads, self.head_dim]), + paddle.reshape(k, [-1, self.kv_num_heads, self.head_dim]), + paddle.reshape(v, [-1, self.kv_num_heads, self.head_dim]), + ) + + def get_kv_cache_shape(self, max_num_blocks: int, kv_cache_quant_type: Optional[str] = None): + if kv_cache_quant_type not in (None, "none"): + raise ValueError("InfLLM-V2 requires an unquantized paged KV cache.") + shape = [max_num_blocks, self.kv_num_heads, self.block_size, self.head_dim] + return shape, shape + + def get_additional_cache_block_bytes(self, cache_dtype_bytes: int) -> int: + if cache_dtype_bytes <= 0: + raise ValueError("cache_dtype_bytes must be positive.") + semantic_slots = self.block_size // self.kernel_stride + semantic_slots += self.block_size // (4 * self.kernel_stride) + return cache_dtype_bytes * self.kv_num_heads * self.head_dim * semantic_slots + + def init_attention_metadata(self, forward_meta: ForwardMeta): + super().init_attention_metadata(forward_meta) + base_metadata = self.attention_metadata + metadata = InfLLMV2AttentionMetadata( + **{field.name: getattr(base_metadata, field.name) for field in fields(FlashAttentionMetadata)} + ) + metadata.compressed_k = self._compressed_k + metadata.compressed_k2 = self._compressed_k2 + metadata.kernel_size = self.kernel_size + metadata.kernel_stride = self.kernel_stride + metadata.topk = self.topk + metadata.dense_len = self.dense_len + metadata.block_size = self.block_size + metadata.init_blocks = self.init_blocks + metadata.local_blocks = self.local_blocks + metadata.selected_capacity = self.selected_capacity + self.attention_metadata = metadata + forward_meta.attn_metadata = metadata + + def reset_runtime_cache(self) -> None: + self._compressed_k = None + self._compressed_k2 = None + self._compressed_cache_owner = None + if hasattr(self, "attention_metadata"): + self.attention_metadata.compressed_k = None + self.attention_metadata.compressed_k2 = None + + def _get_layer_cache(self, layer: paddle.nn.Layer, forward_meta: ForwardMeta): + if layer is None or not hasattr(layer, "layer_id"): + raise ValueError("InfLLM-V2 requires an attention layer with layer_id.") + if getattr(layer, "cache_quant_type_str", "none") != "none": + raise ValueError("InfLLM-V2 requires an unquantized paged KV cache.") + cache_index = 2 * layer.layer_id + caches = forward_meta.caches + if caches is None or len(caches) <= cache_index + 1: + raise ValueError("InfLLM-V2 forward metadata does not contain this layer's K/V cache.") + cache_k, cache_v = caches[cache_index], caches[cache_index + 1] + if cache_k is None or cache_v is None: + raise ValueError("InfLLM-V2 layer K/V cache must not be None.") + return cache_k, cache_v + + def _ensure_compressed_cache(self, cache_k: paddle.Tensor): + owner = self._compressed_cache_owner() if self._compressed_cache_owner is not None else None + if owner is not None and owner._is_shared_buffer_with(cache_k): + return self._compressed_k, self._compressed_k2 + fine_shape = [ + cache_k.shape[0], + self.kv_num_heads, + self.block_size // self.kernel_stride, + self.head_dim, + ] + coarse_shape = [ + cache_k.shape[0], + self.kv_num_heads, + self.block_size // (4 * self.kernel_stride), + self.head_dim, + ] + self._compressed_k = paddle.zeros(fine_shape, dtype=cache_k.dtype) + self._compressed_k2 = paddle.zeros(coarse_shape, dtype=cache_k.dtype) + self._compressed_cache_owner = weakref.ref(cache_k) + self.attention_metadata.compressed_k = self._compressed_k + self.attention_metadata.compressed_k2 = self._compressed_k2 + return self._compressed_k, self._compressed_k2 + + def _ensure_workspace( + self, + query_tokens: int, + max_blocks_per_seq: int, + dtype, + stage2_blocks_per_split: int = _STAGE2_BLOCKS_PER_SPLIT, + allocate_attention: bool = True, + ) -> None: + if max_blocks_per_seq > _STAGE1_MAX_CANDIDATE_BLOCKS: + raise ValueError("InfLLM-V2 Stage 1 supports at most 2048 candidate blocks.") + if stage2_blocks_per_split <= 0: + raise ValueError("InfLLM-V2 Stage 2 blocks per split must be positive.") + key = (query_tokens, max_blocks_per_seq, dtype, stage2_blocks_per_split, allocate_attention) + if key == self._workspace_key: + return + self._topk_indices_ws = paddle.empty([query_tokens, self.kv_num_heads, self.selected_capacity], dtype="int32") + self._block_scores_ws = paddle.empty([query_tokens, self.kv_num_heads, max_blocks_per_seq], dtype="float32") + self._selected_counts_ws = paddle.empty([query_tokens, self.kv_num_heads], dtype="int32") + self._coarse_lse_ws = paddle.empty([query_tokens, self.num_heads], dtype="float32") + max_visible_length = max_blocks_per_seq * self.block_size + coarse_kernel = 4 * self.kernel_size + coarse_stride = 4 * self.kernel_stride + max_coarse_windows = ( + 0 if max_visible_length < coarse_kernel else (max_visible_length - coarse_kernel) // coarse_stride + 1 + ) + coarse_splits = max( + 1, + (max_coarse_windows + _STAGE1_COARSE_WINDOWS_PER_SPLIT - 1) // _STAGE1_COARSE_WINDOWS_PER_SPLIT, + ) + partial_shape = [query_tokens, self.num_heads, coarse_splits] + self._coarse_partial_max_ws = paddle.empty(partial_shape, dtype="float32") + self._coarse_partial_sum_ws = paddle.empty(partial_shape, dtype="float32") + if allocate_attention: + self._attention_out_ws = paddle.empty([query_tokens, self.num_heads, self.head_dim], dtype=dtype) + kv_splits = (self.selected_capacity + stage2_blocks_per_split - 1) // stage2_blocks_per_split + self._partial_acc_ws = paddle.empty( + [query_tokens, self.num_heads, kv_splits, self.head_dim], dtype="float32" + ) + self._partial_max_ws = paddle.empty([query_tokens, self.num_heads, kv_splits], dtype="float32") + self._partial_sum_ws = paddle.empty([query_tokens, self.num_heads, kv_splits], dtype="float32") + else: + self._attention_out_ws = None + self._partial_acc_ws = None + self._partial_max_ws = None + self._partial_sum_ws = None + self._workspace_key = key + + def _release_workspace(self) -> None: + self._workspace_key = None + self._topk_indices_ws = None + self._block_scores_ws = None + self._selected_counts_ws = None + self._coarse_lse_ws = None + self._coarse_partial_max_ws = None + self._coarse_partial_sum_ws = None + self._attention_out_ws = None + self._partial_acc_ws = None + self._partial_max_ws = None + self._partial_sum_ws = None + + @staticmethod + def _load_sparse_ops(): + try: + from fastdeploy.model_executor.ops.gpu import ( + infllmv2_attention_forward, + infllmv2_select_blocks, + infllmv2_update_compressed_k, + ) + except ImportError as exc: + raise RuntimeError( + "INFLLMV2_ATTN requires rebuilt infllmv2_update_compressed_k, " + "infllmv2_select_blocks, and infllmv2_attention_forward custom ops." + ) from exc + return infllmv2_update_compressed_k, infllmv2_select_blocks, infllmv2_attention_forward + + def _update_compressed_cache( + self, current_tokens: paddle.Tensor, cache_k: paddle.Tensor, forward_meta: ForwardMeta + ) -> None: + compressed_k, compressed_k2 = self._ensure_compressed_cache(cache_k) + update_compressed_k, _, _ = self._load_sparse_ops() + outputs = update_compressed_k( + current_tokens, + cache_k, + compressed_k, + compressed_k2, + forward_meta.block_tables, + forward_meta.seq_lens_decoder, + forward_meta.seq_lens_this_time, + forward_meta.batch_id_per_token, + forward_meta.cu_seqlens_q, + self.kernel_size, + self.kernel_stride, + ) + if isinstance(outputs, (tuple, list)): + self._compressed_k, self._compressed_k2 = outputs + self.attention_metadata.compressed_k = self._compressed_k + self.attention_metadata.compressed_k2 = self._compressed_k2 + + def _prepare_sparse_runtime(self, layer: paddle.nn.Layer, forward_meta: ForwardMeta) -> None: + metadata = self.attention_metadata + if self.pd_disaggregation_mode == "per_query": + metadata.kv_signal_data_list[layer.layer_id] = init_signal_layerwise( + metadata.kv_signal_metadata, layer.layer_id + self.start_layer_index + ) + if int(os.getenv("USE_TBO", "0")) == 1 and hasattr(forward_meta, "tbo_microbatch_id"): + os.environ["FLAGS_fmt_write_cache_completed_signal"] = str(forward_meta.tbo_microbatch_id) + if layer.layer_id == 0: + get_block_shape_and_split_kv_block( + forward_meta.seq_lens_encoder, + forward_meta.seq_lens_decoder, + forward_meta.seq_lens_this_time, + forward_meta.decoder_batch_ids, + forward_meta.decoder_tile_ids_per_batch, + forward_meta.decoder_num_blocks_cpu, + forward_meta.decoder_num_blocks_device, + forward_meta.decoder_chunk_size_device, + forward_meta.max_len_tensor_cpu, + forward_meta.encoder_batch_ids, + forward_meta.encoder_tile_ids_per_batch, + forward_meta.encoder_num_blocks_x_cpu, + forward_meta.kv_batch_ids, + forward_meta.kv_tile_ids_per_batch, + forward_meta.kv_num_blocks_x_cpu, + self.encoder_block_shape_q, + self.decoder_block_shape_q, + self.group_size, + self.block_size, + ) + + def _prepare_sparse_prefill_runtime(self, layer: paddle.nn.Layer, forward_meta: ForwardMeta) -> None: + """Prepare the model-specific cache metadata required by sparse prefill.""" + self._prepare_sparse_runtime(layer, forward_meta) + if layer.layer_id != 0 or forward_meta.max_len_tensor_cpu[1].item() <= 0: + return + ( + forward_meta.cu_seqlens_k, + forward_meta.pre_cache_batch_ids, + forward_meta.pre_cache_tile_ids_per_batch, + forward_meta.pre_cache_num_blocks_cpu, + forward_meta.kv_token_num_cpu, + ) = pre_cache_len_concat( + forward_meta.seq_lens_encoder, + forward_meta.seq_lens_decoder, + forward_meta.seq_lens_this_time, + forward_meta.max_len_tensor_cpu[2], + self.block_size, + ) + + def _write_prefill_cache( + self, + qkv: paddle.Tensor, + cache_k: paddle.Tensor, + cache_v: paddle.Tensor, + layer: paddle.nn.Layer, + forward_meta: ForwardMeta, + ): + norm_after_rope = not getattr(layer, "qk_norm_before_rope", False) + q_norm_weight = getattr(layer, "q_norm_weight", None) if norm_after_rope else None + k_norm_weight = getattr(layer, "k_norm_weight", None) if norm_after_rope else None + return gqa_rope_write_cache( + qkv, + cache_k, + cache_v, + forward_meta.cu_seqlens_q, + forward_meta.cu_seqlens_k, + forward_meta.rotary_embs, + forward_meta.seq_lens_this_time, + forward_meta.seq_lens_encoder, + forward_meta.seq_lens_decoder, + forward_meta.batch_id_per_token, + forward_meta.block_tables, + forward_meta.kv_batch_ids, + forward_meta.kv_tile_ids_per_batch, + forward_meta.kv_num_blocks_x_cpu, + forward_meta.pre_cache_batch_ids, + forward_meta.pre_cache_tile_ids_per_batch, + forward_meta.pre_cache_num_blocks_cpu, + q_norm_weight, + k_norm_weight, + None, + None, + getattr(layer, "cache_k_out_scale", None), + getattr(layer, "cache_v_out_scale", None), + getattr(layer, "cache_k_zp", None), + getattr(layer, "cache_v_zp", None), + self.attention_metadata.kv_signal_data_list[layer.layer_id], + forward_meta.kv_token_num_cpu[0].item(), + self.max_seq_len, + getattr(layer, "rms_norm_eps", 1e-6), + layer.use_neox_rotary_style, + "none", + self.rope_3d, + )[:3] + + def _write_decode_cache( + self, + qkv: paddle.Tensor, + cache_k: paddle.Tensor, + cache_v: paddle.Tensor, + layer: paddle.nn.Layer, + forward_meta: ForwardMeta, + ) -> paddle.Tensor: + norm_after_rope = not getattr(layer, "qk_norm_before_rope", False) + q_norm_weight = getattr(layer, "q_norm_weight", None) if norm_after_rope else None + k_norm_weight = getattr(layer, "k_norm_weight", None) if norm_after_rope else None + return decoder_write_cache_with_rope( + qkv, + cache_k, + cache_v, + forward_meta.seq_lens_encoder, + forward_meta.seq_lens_decoder, + forward_meta.seq_lens_this_time, + forward_meta.batch_id_per_token, + forward_meta.cu_seqlens_q, + forward_meta.block_tables, + forward_meta.max_len_tensor_cpu, + forward_meta.rotary_embs, + layer.qkv_bias, + getattr(layer, "cache_k_scale", None), + getattr(layer, "cache_v_scale", None), + getattr(layer, "cache_k_out_scale", None), + getattr(layer, "cache_v_out_scale", None), + getattr(layer, "cache_k_zp", None), + getattr(layer, "cache_v_zp", None), + self.attention_metadata.kv_signal_data_list[layer.layer_id], + q_norm_weight, + k_norm_weight, + getattr(layer, "rms_norm_eps", 1e-6), + "none", + layer.use_neox_rotary_style, + self.rope_3d, + self.max_seq_len, + getattr(layer, "quant_max_bound", 0.0), + getattr(layer, "quant_min_bound", 0.0), + False, + ) + + def _trace_rank_matches(self) -> bool: + rank = int(os.getenv("PADDLE_TRAINER_ID", os.getenv("RANK", "0"))) + return rank == self._selector_trace_rank + + def _trace_request_slot(self, batch_id: int, query_position: int) -> int: + last_position = self._selector_trace_last_position.get(batch_id) + if last_position is None or query_position <= last_position: + request_index = self._selector_trace_next_request_index + self._selector_trace_next_request_index += 1 + self._selector_trace_request_index[batch_id] = request_index + self._selector_trace_last_position[batch_id] = query_position + return self._selector_trace_request_index[batch_id] + + def _write_selector_trace(self) -> None: + payload = { + "schema_version": 1, + "kind": "infllmv2_selector_samples", + "reference_definition": _SELECTOR_REFERENCE_DEFINITION, + "rank": self._selector_trace_rank, + "layer": self._selector_trace_layer, + "block_size": self.block_size, + "kernel_size": self.kernel_size, + "kernel_stride": self.kernel_stride, + "topk": self.topk, + "dense_len": self.dense_len, + "init_blocks": self.init_blocks, + "local_blocks": self.local_blocks, + "selected_capacity": self.selected_capacity, + "samples": self._selector_trace_samples, + } + temporary_path = self._selector_trace_path.with_name(f".{self._selector_trace_path.name}.tmp") + with temporary_path.open("w", encoding="utf-8") as trace_file: + json.dump(payload, trace_file, indent=2, sort_keys=True) + trace_file.write("\n") + os.replace(temporary_path, self._selector_trace_path) + + @paddle.no_grad() + def _record_selector_trace( + self, + post_rope_q: paddle.Tensor, + cache_k: paddle.Tensor, + forward_meta: ForwardMeta, + layer: paddle.nn.Layer, + topk_indices: paddle.Tensor, + selected_counts: paddle.Tensor, + ) -> None: + if ( + self._selector_trace_path is None + or layer.layer_id != self._selector_trace_layer + or not self._trace_rank_matches() + or forward_meta.is_dummy_or_profile_run + or len(self._selector_trace_samples) >= self._selector_trace_max_samples + ): + return + + batch_ids = forward_meta.batch_id_per_token.numpy().tolist() + seq_lens_decoder = forward_meta.seq_lens_decoder.numpy().tolist() + cu_seqlens_q = forward_meta.cu_seqlens_q.numpy().tolist() + counts = selected_counts.numpy() + selected = topk_indices.numpy() + group_size = self.num_heads // self.kv_num_heads + attention_scale = 1.0 / math.sqrt(self.head_dim) + + for token_id, batch_id in enumerate(batch_ids): + if len(self._selector_trace_samples) >= self._selector_trace_max_samples: + break + if batch_id < 0 or batch_id >= len(seq_lens_decoder): + raise ValueError(f"InfLLM-V2 selector trace received invalid batch id {batch_id}.") + query_offset = token_id - cu_seqlens_q[batch_id] + query_position = seq_lens_decoder[batch_id] + query_offset + visible_length = query_position + 1 + if visible_length < self.dense_len: + continue + valid_blocks = (visible_length + self.block_size - 1) // self.block_size + physical_blocks = forward_meta.block_tables[batch_id, :valid_blocks].astype("int64") + logical_k = paddle.index_select(cache_k, physical_blocks, axis=0) + logical_k = paddle.transpose(logical_k, [1, 0, 2, 3]) + logical_k = paddle.reshape(logical_k, [self.kv_num_heads, valid_blocks * self.block_size, self.head_dim]) + logical_k = paddle.cast(logical_k[:, :visible_length, :], "float32") + grouped_q = paddle.reshape( + paddle.cast(post_rope_q[token_id], "float32"), + [self.kv_num_heads, group_size, self.head_dim], + ) + logits = paddle.matmul(grouped_q, paddle.transpose(logical_k, [0, 2, 1])) * attention_scale + probabilities = paddle.nn.functional.softmax(logits, axis=-1) + padded_tokens = valid_blocks * self.block_size - visible_length + if padded_tokens: + probabilities = paddle.nn.functional.pad(probabilities, [0, padded_tokens]) + block_mass = paddle.reshape( + probabilities, + [self.kv_num_heads, group_size, valid_blocks, self.block_size], + ).sum(axis=[1, 3]) + request_index = self._trace_request_slot(batch_id, query_position) + + for kv_head in range(self.kv_num_heads): + if len(self._selector_trace_samples) >= self._selector_trace_max_samples: + break + selected_count = int(counts[token_id, kv_head]) + if selected_count <= 0 or selected_count > valid_blocks: + raise ValueError( + f"InfLLM-V2 selector trace selected count {selected_count} is invalid for {valid_blocks} blocks." + ) + selected_blocks = sorted(int(value) for value in selected[token_id, kv_head, :selected_count]) + reference = paddle.topk(block_mass[kv_head], k=selected_count, largest=True, sorted=False).indices + reference_blocks = sorted(int(value) for value in reference.numpy().tolist()) + query_head_start = kv_head * group_size + self._selector_trace_samples.append( + { + "rank": self._selector_trace_rank, + "layer": self._selector_trace_layer, + "request_index": request_index, + "query_index": token_id, + "query_offset": query_offset, + "query_position": query_position, + "kv_head": kv_head, + "query_head_start": query_head_start, + "query_head_end": query_head_start + group_size, + "block_size": self.block_size, + "topk": self.topk, + "selected_count": selected_count, + "selected_blocks": selected_blocks, + "reference_blocks": reference_blocks, + "reference_metric": _SELECTOR_REFERENCE_DEFINITION, + } + ) + if self._selector_trace_samples: + self._write_selector_trace() + + def _can_use_sparse_prefill(self, qkv: paddle.Tensor, forward_meta: ForwardMeta) -> bool: + """Return whether this forward has the single-request causal layout supported by sparse prefill.""" + if not self.sparse_prefill or not self.causal or qkv is None: + return False + if forward_meta.block_tables.shape[0] != 1: + return False + if getattr(forward_meta, "attn_mask_offsets", None) is not None: + return False + if forward_meta.max_len_tensor_cpu[2].item() != 0: + # A non-zero decoder prefix is the shared chunked-prefill path. It + # keeps the dense fallback until that path can supply block-aligned + # per-chunk metadata without a device-to-host synchronization. + return False + prompt_tokens = int(qkv.shape[0]) + sparse_start = ((self.dense_len + self.block_size - 1) // self.block_size) * self.block_size + selected_count = self.topk + self.local_blocks + current_blocks = (_PREFILL_QUERY_TILE_SIZE + self.block_size - 1) // self.block_size + has_full_selection = selected_count > current_blocks and sparse_start // self.block_size + 1 >= selected_count + production_shape = ( + self.block_size == 64 + and self.head_dim == 128 + and self.num_heads // self.kv_num_heads == 16 + and qkv.dtype in (paddle.float16, paddle.bfloat16) + ) + return production_shape and has_full_selection and prompt_tokens > sparse_start + + @staticmethod + def _single_sequence_cu_seqlens(length: int, place) -> paddle.Tensor: + return paddle.to_tensor([0, length], dtype="int32", place=place) + + def _dense_prefill_segment( + self, + query: paddle.Tensor, + key: paddle.Tensor, + value: paddle.Tensor, + ) -> paddle.Tensor: + query_length = int(query.shape[0]) + key_length = int(key.shape[0]) + cu_query = self._single_sequence_cu_seqlens(query_length, query.place) + cu_key = self._single_sequence_cu_seqlens(key_length, query.place) + return flash_attn_func( + query, + key, + value, + cu_query, + cu_key, + max_seqlen_q=query_length, + max_seqlen_k=key_length, + causal=True, + num_heads=self.num_heads, + kv_num_heads=self.kv_num_heads, + head_dim=self.head_dim, + version=2, + )[0] + + def _select_prefill_query_blocks( + self, + query: paddle.Tensor, + block_tables: paddle.Tensor, + sparse_start: int, + sparse_end: int, + ) -> paddle.Tensor: + sparse_tokens = sparse_end - sparse_start + full_tiles = sparse_tokens // _PREFILL_QUERY_TILE_SIZE + remainder = sparse_tokens % _PREFILL_QUERY_TILE_SIZE + representative_parts = [] + position_parts = [] + if full_tiles: + full_end = sparse_start + full_tiles * _PREFILL_QUERY_TILE_SIZE + representative_parts.append( + query[sparse_start + _PREFILL_QUERY_TILE_SIZE - 1 : full_end : _PREFILL_QUERY_TILE_SIZE] + ) + position_parts.append( + paddle.arange( + sparse_start + _PREFILL_QUERY_TILE_SIZE - 1, + full_end, + _PREFILL_QUERY_TILE_SIZE, + dtype="int32", + ) + ) + if remainder: + representative_parts.append(query[sparse_end - 1 : sparse_end]) + position_parts.append(paddle.to_tensor([sparse_end - 1], dtype="int32", place=query.place)) + representative_query = ( + representative_parts[0] if len(representative_parts) == 1 else paddle.concat(representative_parts, axis=0) + ) + query_tiles = int(representative_query.shape[0]) + pseudo_block_tables = paddle.tile(block_tables[:1], [query_tiles, 1]) + pseudo_seq_lens_decoder = ( + position_parts[0] if len(position_parts) == 1 else paddle.concat(position_parts, axis=0) + ) + pseudo_seq_lens_this_time = paddle.ones([query_tiles], dtype="int32") + pseudo_batch_ids = paddle.arange(query_tiles, dtype="int32") + pseudo_cu_seqlens_q = paddle.arange(query_tiles + 1, dtype="int32") + + self._ensure_workspace( + query_tiles, + int(block_tables.shape[1]), + query.dtype, + allocate_attention=False, + ) + _, select_blocks, _ = self._load_sparse_ops() + selection = select_blocks( + representative_query, + self._compressed_k, + self._compressed_k2, + pseudo_block_tables, + pseudo_seq_lens_decoder, + pseudo_seq_lens_this_time, + pseudo_batch_ids, + pseudo_cu_seqlens_q, + self._topk_indices_ws, + self._block_scores_ws, + self._selected_counts_ws, + self._coarse_lse_ws, + self._coarse_partial_max_ws, + self._coarse_partial_sum_ws, + self.block_size, + self.kernel_size, + self.kernel_stride, + self.topk, + self.dense_len, + self.init_blocks, + self.local_blocks, + ) + return selection[0] + + def _gather_prefill_cache_blocks( + self, + cache: paddle.Tensor, + block_tables: paddle.Tensor, + logical_blocks: paddle.Tensor, + ) -> paddle.Tensor: + query_tiles = int(logical_blocks.shape[0]) + selected_blocks = int(logical_blocks.shape[2]) + expanded_tables = paddle.tile(block_tables[:1].unsqueeze(1), [query_tiles, self.kv_num_heads, 1]) + physical_blocks = paddle.take_along_axis(expanded_tables, logical_blocks, axis=2) + kv_heads = paddle.arange(self.kv_num_heads, dtype="int32").reshape([1, -1, 1]) + kv_heads = paddle.tile(kv_heads, [query_tiles, 1, selected_blocks]) + gather_indices = paddle.stack([physical_blocks, kv_heads], axis=-1) + gathered = paddle.gather_nd(cache, gather_indices) + return paddle.transpose(gathered, [0, 2, 3, 1, 4]).reshape( + [query_tiles, selected_blocks * self.block_size, self.kv_num_heads, self.head_dim] + ) + + def _sparse_prefill_tile_batch( + self, + query: paddle.Tensor, + key_cache: paddle.Tensor, + value_cache: paddle.Tensor, + block_tables: paddle.Tensor, + selected_blocks: paddle.Tensor, + first_position: int, + query_tile_size: int, + ) -> paddle.Tensor: + query_tiles = int(selected_blocks.shape[0]) + current_blocks = (query_tile_size + self.block_size - 1) // self.block_size + selected_count = self.topk + self.local_blocks + history_count = selected_count - current_blocks + if history_count <= 0: + raise ValueError("InfLLM-V2 sparse prefill requires at least one selected history block.") + + # Stage 1 returns logical block IDs in ascending order and forces the + # current logical blocks into the selection, so those blocks occupy the + # final slots and the preceding slots are fully visible history. + history = selected_blocks[:, :, :history_count] + history_k = self._gather_prefill_cache_blocks(key_cache, block_tables, history) + history_v = self._gather_prefill_cache_blocks(value_cache, block_tables, history) + + first_logical_block = first_position // self.block_size + current = paddle.arange( + first_logical_block, + first_logical_block + query_tiles * current_blocks, + dtype="int32", + ).reshape([query_tiles, 1, current_blocks]) + current = paddle.tile(current, [1, self.kv_num_heads, 1]) + current_k = self._gather_prefill_cache_blocks(key_cache, block_tables, current) + current_v = self._gather_prefill_cache_blocks(value_cache, block_tables, current) + current_k = current_k[:, :query_tile_size] + current_v = current_v[:, :query_tile_size] + + tiled_query = query.reshape([query_tiles, query_tile_size, self.num_heads, self.head_dim]) + history_output, _, history_lse, _ = paddle._C_ops.flash_attn( + tiled_query, + history_k, + history_v, + None, + None, + 0.0, + False, + False, + True, + "", + ) + current_output, _, current_lse, _ = paddle._C_ops.flash_attn( + tiled_query, + current_k, + current_v, + None, + None, + 0.0, + True, + False, + True, + "", + ) + history_lse = paddle.transpose(history_lse, [0, 2, 1]).unsqueeze(-1) + current_lse = paddle.transpose(current_lse, [0, 2, 1]).unsqueeze(-1) + maximum_lse = paddle.maximum(history_lse, current_lse) + history_weight = paddle.exp(history_lse - maximum_lse).astype(query.dtype) + current_weight = paddle.exp(current_lse - maximum_lse).astype(query.dtype) + output = (history_output * history_weight + current_output * current_weight) / ( + history_weight + current_weight + ) + return output.reshape([-1, self.num_heads, self.head_dim]) + + def _sparse_prefill_attention( + self, + query: paddle.Tensor, + key: paddle.Tensor, + value: paddle.Tensor, + key_cache: paddle.Tensor, + value_cache: paddle.Tensor, + block_tables: paddle.Tensor, + ) -> paddle.Tensor: + total_tokens = int(query.shape[0]) + sparse_start = ((self.dense_len + self.block_size - 1) // self.block_size) * self.block_size + sparse_end = total_tokens + outputs = [self._dense_prefill_segment(query[:sparse_start], key[:sparse_start], value[:sparse_start])] + selected_blocks = self._select_prefill_query_blocks(query, block_tables, sparse_start, sparse_end) + sparse_tokens = sparse_end - sparse_start + full_tiles = sparse_tokens // _PREFILL_QUERY_TILE_SIZE + tiles_per_chunk = self.prefill_query_chunk_size // _PREFILL_QUERY_TILE_SIZE + for first_tile in range(0, full_tiles, tiles_per_chunk): + last_tile = min(full_tiles, first_tile + tiles_per_chunk) + token_begin = sparse_start + first_tile * _PREFILL_QUERY_TILE_SIZE + token_end = sparse_start + last_tile * _PREFILL_QUERY_TILE_SIZE + outputs.append( + self._sparse_prefill_tile_batch( + query[token_begin:token_end], + key_cache, + value_cache, + block_tables, + selected_blocks[first_tile:last_tile], + token_begin, + _PREFILL_QUERY_TILE_SIZE, + ) + ) + remainder = sparse_tokens % _PREFILL_QUERY_TILE_SIZE + if remainder: + token_begin = sparse_start + full_tiles * _PREFILL_QUERY_TILE_SIZE + outputs.append( + self._sparse_prefill_tile_batch( + query[token_begin:sparse_end], + key_cache, + value_cache, + block_tables, + selected_blocks[full_tiles : full_tiles + 1], + token_begin, + remainder, + ) + ) + return paddle.concat(outputs, axis=0) + + def _forward_sparse_prefill( + self, qkv: paddle.Tensor, layer: paddle.nn.Layer, forward_meta: ForwardMeta + ) -> paddle.Tensor: + cache_k, cache_v = self._get_layer_cache(layer, forward_meta) + self._prepare_sparse_prefill_runtime(layer, forward_meta) + post_rope_q, post_rope_k, post_rope_v = self._write_prefill_cache(qkv, cache_k, cache_v, layer, forward_meta) + self._update_compressed_cache(qkv, cache_k, forward_meta) + try: + output = self._sparse_prefill_attention( + post_rope_q, + post_rope_k, + post_rope_v, + cache_k, + cache_v, + forward_meta.block_tables, + ) + finally: + self._release_workspace() + if layer.layer_id == 0 and not self._logged_sparse_prefill_activation: + logger.info( + "InfLLM-V2 sparse prefill activated: " + f"query_tile={_PREFILL_QUERY_TILE_SIZE}, query_chunk={self.prefill_query_chunk_size}, " + f"selected_blocks={self.topk + self.local_blocks}, dense_len={self.dense_len}" + ) + self._logged_sparse_prefill_activation = True + return paddle.reshape(output, [-1, self.num_heads * self.head_dim]) + + def _forward_sparse_decode( + self, qkv: paddle.Tensor, layer: paddle.nn.Layer, forward_meta: ForwardMeta + ) -> paddle.Tensor: + cache_k, cache_v = self._get_layer_cache(layer, forward_meta) + owner = self._compressed_cache_owner() if self._compressed_cache_owner is not None else None + if owner is None or not owner._is_shared_buffer_with(cache_k): + raise RuntimeError( + "InfLLM-V2 sparse decode requires semantic K summaries initialized " + "by prefill on the same paged KV cache." + ) + self._prepare_sparse_runtime(layer, forward_meta) + post_rope_qkv = self._write_decode_cache(qkv, cache_k, cache_v, layer, forward_meta) + post_rope_q, _, _ = self._split_qkv(post_rope_qkv) + self._update_compressed_cache(post_rope_qkv, cache_k, forward_meta) + + _, select_blocks, sparse_attention = self._load_sparse_ops() + self._ensure_workspace(int(post_rope_q.shape[0]), int(forward_meta.block_tables.shape[1]), post_rope_q.dtype) + selection = select_blocks( + post_rope_q, + self._compressed_k, + self._compressed_k2, + forward_meta.block_tables, + forward_meta.seq_lens_decoder, + forward_meta.seq_lens_this_time, + forward_meta.batch_id_per_token, + forward_meta.cu_seqlens_q, + self._topk_indices_ws, + self._block_scores_ws, + self._selected_counts_ws, + self._coarse_lse_ws, + self._coarse_partial_max_ws, + self._coarse_partial_sum_ws, + self.block_size, + self.kernel_size, + self.kernel_stride, + self.topk, + self.dense_len, + self.init_blocks, + self.local_blocks, + ) + metadata = self.attention_metadata + metadata.topk_indices, metadata.block_scores, metadata.selected_counts = selection[:3] + self._record_selector_trace( + post_rope_q, + cache_k, + forward_meta, + layer, + metadata.topk_indices, + metadata.selected_counts, + ) + output = sparse_attention( + post_rope_q, + cache_k, + cache_v, + forward_meta.block_tables, + forward_meta.seq_lens_decoder, + forward_meta.seq_lens_this_time, + forward_meta.batch_id_per_token, + forward_meta.cu_seqlens_q, + metadata.topk_indices, + self._attention_out_ws, + self._partial_acc_ws, + self._partial_max_ws, + self._partial_sum_ws, + ) + if isinstance(output, (tuple, list)): + output = output[0] + if layer.layer_id == 0 and not self._logged_sparse_activation: + logger.info( + "InfLLM-V2 sparse decode activated: " + f"block_size={self.block_size}, kernel={self.kernel_size}/{self.kernel_stride}, " + f"topk={self.topk}, local_blocks={self.local_blocks}, dense_len={self.dense_len}" + ) + self._logged_sparse_activation = True + return paddle.reshape(output, [-1, self.num_heads * self.head_dim]) + + def _forward_dense_and_update( + self, + q: paddle.Tensor, + k: paddle.Tensor, + v: paddle.Tensor, + qkv: paddle.Tensor, + compressed_kv: paddle.Tensor, + k_pe: paddle.Tensor, + layer: paddle.nn.Layer, + forward_meta: ForwardMeta, + ) -> paddle.Tensor: + output = super().forward_mixed(q, k, v, qkv, compressed_kv, k_pe, layer, forward_meta) + cache_k, _ = self._get_layer_cache(layer, forward_meta) + current_tokens = qkv if qkv is not None else q + if current_tokens is None: + raise ValueError("InfLLM-V2 prefill requires qkv or q tokens for semantic-cache update.") + self._update_compressed_cache(current_tokens, cache_k, forward_meta) + return output + + def forward_decode(self, q, k, v, qkv, compressed_kv, k_pe, layer, forward_meta): + return self._forward_sparse_decode(qkv, layer, forward_meta) + + def forward_extend(self, q, k, v, qkv, compressed_kv, k_pe, layer, forward_meta): + if self._can_use_sparse_prefill(qkv, forward_meta): + return self._forward_sparse_prefill(qkv, layer, forward_meta) + return self._forward_dense_and_update(q, k, v, qkv, compressed_kv, k_pe, layer, forward_meta) + + def forward_mixed(self, q, k, v, qkv, compressed_kv, k_pe, layer, forward_meta): + if getattr(forward_meta, "exist_prefill", False): + if self._can_use_sparse_prefill(qkv, forward_meta): + return self._forward_sparse_prefill(qkv, layer, forward_meta) + return self._forward_dense_and_update(q, k, v, qkv, compressed_kv, k_pe, layer, forward_meta) + return self._forward_sparse_decode(qkv, layer, forward_meta) + + def forward_native_backend(self, q, k, v, qkv, layer, forward_meta): + return self._forward_dense_and_update(q, k, v, qkv, None, None, layer, forward_meta) diff --git a/fastdeploy/model_executor/layers/rotary_embedding.py b/fastdeploy/model_executor/layers/rotary_embedding.py index b9eb9ac5bd7..00da3022a5f 100644 --- a/fastdeploy/model_executor/layers/rotary_embedding.py +++ b/fastdeploy/model_executor/layers/rotary_embedding.py @@ -138,6 +138,48 @@ def __call__(self, position_ids): return rot_emb +class MiniCPMLongRotaryEmbedding: + """MiniCPM LongRoPE with the full-dimension NeoX cache layout.""" + + def __init__( + self, + rotary_dim: int, + base: float, + max_position_embeddings: int, + short_factor: list[float], + long_factor: list[float], + original_max_position_embeddings: int, + ) -> None: + factor_count = rotary_dim // 2 + if len(short_factor) != factor_count or len(long_factor) != factor_count: + raise ValueError(f"MiniCPM LongRoPE factors must contain {factor_count} values") + if min(short_factor) <= 0 or min(long_factor) <= 0: + raise ValueError("MiniCPM LongRoPE factors must be positive") + + self.rotary_dim = rotary_dim + self.base = base + self.short_factor = short_factor + self.long_factor = long_factor + self.original_max_position_embeddings = original_max_position_embeddings + scale = max_position_embeddings / original_max_position_embeddings + self.magnitude_scale = math.sqrt(1 + math.log(scale) / math.log(original_max_position_embeddings)) + + def __call__(self, position_ids: paddle.Tensor) -> paddle.Tensor: + bsz, max_seq_len = position_ids.shape[:2] + factors = self.long_factor if max_seq_len > self.original_max_position_embeddings else self.short_factor + inv_freq = self.base ** (-paddle.arange(0, self.rotary_dim, 2, dtype="float32") / self.rotary_dim) + inv_freq = inv_freq / paddle.to_tensor(factors, dtype="float32") + # This is an outer product, so avoid routing the small startup-only + # calculation through cuBLAS. + freqs = position_ids.cast("float32").unsqueeze(-1) * inv_freq.reshape([1, 1, -1]) + emb = paddle.concat([freqs, freqs], axis=-1).reshape((bsz, max_seq_len, 1, self.rotary_dim)) + rotary_emb = paddle.stack( + [paddle.cos(emb) * self.magnitude_scale, paddle.sin(emb) * self.magnitude_scale], axis=0 + ) + rotary_emb.stop_gradient = True + return rotary_emb + + def yarn_get_mscale(scale=1, mscale=1): """ """ if scale <= 1: @@ -376,6 +418,22 @@ def get_rope_impl( use_neox_rotary_style=True, ) rotary_emb = rotary_emb_layer(position_ids) + elif architecture == "MiniCPMForCausalLM": + rope_scaling = getattr(model_config, "rope_scaling", None) + if rope_scaling is None: + rotary_emb_layer = QwenRotaryEmbedding(rotary_dim, base, partial_rotary_factor) + elif rope_scaling.get("rope_type", rope_scaling.get("type", "")) == "longrope": + rotary_emb_layer = MiniCPMLongRotaryEmbedding( + rotary_dim=rotary_dim, + base=base, + max_position_embeddings=model_config.max_position_embeddings, + short_factor=rope_scaling["short_factor"], + long_factor=rope_scaling["long_factor"], + original_max_position_embeddings=rope_scaling["original_max_position_embeddings"], + ) + else: + raise ValueError("MiniCPMForCausalLM requires standard RoPE or longrope") + rotary_emb = rotary_emb_layer(position_ids) else: rotary_emb_layer = ErnieRotaryEmbedding(rotary_dim, base, partial_rotary_factor) rotary_emb = rotary_emb_layer(position_ids) diff --git a/fastdeploy/model_executor/logits_processor/thinking_budget.py b/fastdeploy/model_executor/logits_processor/thinking_budget.py index b8493c1756d..45ae53d0727 100644 --- a/fastdeploy/model_executor/logits_processor/thinking_budget.py +++ b/fastdeploy/model_executor/logits_processor/thinking_budget.py @@ -16,13 +16,13 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Dict, Optional import paddle from paddleformers.utils.log import logger -from fastdeploy.config import FDConfig +from fastdeploy.config import FDConfig, validate_thinking_token_sequences from fastdeploy.model_executor.logits_processor.base import LogitsProcessor @@ -37,6 +37,11 @@ class _ThinkingState: stop_sentence_token_ids: Optional[list[int]] = None stop_sentence_pos: int = 0 prompt_checked: bool = False + # Multi-token marker mode: partial-match buffers and forced-emission progress + start_buffer: list[int] = field(default_factory=list) + end_buffer: list[int] = field(default_factory=list) + forcing: bool = False + forcing_pos: int = 0 class ThinkingBudgetLogitsProcessor(LogitsProcessor): @@ -50,7 +55,11 @@ class ThinkingBudgetLogitsProcessor(LogitsProcessor): {"thinking_budget": } Requires model_config to provide think_start_id and think_end_id. If any of - these are missing or invalid (-1), the processor will be disabled. + these are missing or invalid (-1), the processor falls back to multi-token + marker sequences from model_config.think_token_sequences (a dict with + required "start"/"end" sequence lists and a "forced_end" id list), which + supports tokenizers where / are not single vocab tokens. + Without either form of markers, the processor is disabled. """ def __init__(self, fd_config: FDConfig) -> None: @@ -62,7 +71,21 @@ def __init__(self, fd_config: FDConfig) -> None: self.think_start_token_id = think_start_id if isinstance(think_start_id, int) and think_start_id >= 0 else -1 self.think_end_token_id = think_end_id if isinstance(think_end_id, int) and think_end_id >= 0 else -1 self.line_break_token_id = line_break_id if isinstance(line_break_id, int) and line_break_id >= 0 else -1 - self._enabled = self.think_start_token_id >= 0 and self.think_end_token_id >= 0 + self._single_token_mode = self.think_start_token_id >= 0 and self.think_end_token_id >= 0 + sequences = validate_thinking_token_sequences( + getattr(fd_config.model_config, "think_token_sequences", None), fd_config.model_config.vocab_size + ) + self.think_start_sequences: tuple[tuple[int, ...], ...] = () + self.think_end_sequences: tuple[tuple[int, ...], ...] = () + self.think_forced_end_ids: list[int] = [] + if sequences is not None: + if self._single_token_mode: + raise ValueError("Configure either single-token thinking markers or think_token_sequences, not both") + self.think_start_sequences = tuple(dict.fromkeys(tuple(sequence) for sequence in sequences["start"])) + self.think_end_sequences = tuple(dict.fromkeys(tuple(sequence) for sequence in sequences["end"])) + self.think_forced_end_ids = list(sequences["forced_end"]) + self._sequence_mode = sequences is not None + self._enabled = self._single_token_mode or self._sequence_mode if not self._enabled: logger.warning( "ThinkingBudgetLogitsProcessor disabled: missing token ids " @@ -74,6 +97,106 @@ def __init__(self, fd_config: FDConfig) -> None: self._active_budgets: list[int] = [] self._active_slots: list[int] = [] + def _resolve_thinking_budget(self, logit_proc_args, slot_id: int, share_inputs: dict) -> Optional[int]: + """Per-request thinking budget hook; model wrappers may override the source.""" + thinking_budget = logit_proc_args.get("thinking_budget") if logit_proc_args else None + if thinking_budget is None or not isinstance(thinking_budget, int) or thinking_budget < 0: + return None + return thinking_budget + + @staticmethod + def _buffer_match(buffer: list[int], sequences: tuple[tuple[int, ...], ...]) -> int: + """2 = buffer exactly matches a sequence, 1 = proper prefix, 0 = mismatch.""" + is_prefix = False + for sequence in sequences: + if list(sequence[: len(buffer)]) == buffer: + if len(sequence) == len(buffer): + return 2 + is_prefix = True + return 1 if is_prefix else 0 + + def _consume_token_sequences(self, state: _ThinkingState, token_id: int) -> None: + """Advance the multi-token marker state machine by one token. + + Tokens held in a prefix buffer are only counted as thinking content once + they are ruled out as marker candidates, mirroring single-token mode + where marker tokens never consume the budget. + """ + if state.forcing: + forced_sequence = (state.stop_sentence_token_ids or []) + self.think_forced_end_ids + expected_token_id = forced_sequence[state.forcing_pos] + if token_id != expected_token_id: + raise RuntimeError( + "Forced thinking-end token mismatch: " f"expected {expected_token_id}, received {token_id}" + ) + state.forcing_pos += 1 + if state.forcing_pos == len(forced_sequence): + state.forcing = False + state.forcing_pos = 0 + state.ended = True + state.start_buffer.clear() + state.end_buffer.clear() + return + if state.started and not state.ended: + state.end_buffer.append(token_id) + while state.end_buffer: + status = self._buffer_match(state.end_buffer, self.think_end_sequences) + if status == 2: + state.ended = True + state.forcing = False + state.forcing_pos = 0 + state.end_buffer.clear() + state.start_buffer.clear() + return + if status == 1: + return + state.end_buffer.pop(0) + state.tokens_after_start += 1 + return + state.start_buffer.append(token_id) + while state.start_buffer: + status = self._buffer_match(state.start_buffer, self.think_start_sequences) + if status == 2: + # Matches in both idle and ended phase, so a new thinking round + # after a natural end restarts the budget. + state.started = True + state.ended = False + state.tokens_after_start = 0 + state.stop_sentence_pos = 0 + state.forcing = False + state.forcing_pos = 0 + state.start_buffer.clear() + state.end_buffer.clear() + return + if status == 1: + return + state.start_buffer.pop(0) + + def _scan_prompt_state_sequences(self, state: _ThinkingState, prompt_slice: list) -> None: + """Replay the prompt through the sequence matcher; prompt tokens never + consume the budget, and a partial marker at the prompt tail stays buffered.""" + for token_id in prompt_slice: + token_id = int(token_id) + if token_id >= 0: + self._consume_token_sequences(state, token_id) + state.tokens_after_start = 0 + if prompt_slice: + state.last_token_id = int(prompt_slice[-1]) + + def _consume_decode_token_sequences( + self, state: _ThinkingState, last_token_id: int, current_step_idx: Optional[int] + ) -> None: + # Without step indices, repeated identical tokens cannot be told apart from + # duplicate update_state calls, so fall back to value-based dedup there. + if current_step_idx is None: + if state.forcing or last_token_id != state.last_token_id: + state.last_token_id = last_token_id + self._consume_token_sequences(state, last_token_id) + elif current_step_idx != state.last_step_idx: + state.last_step_idx = current_step_idx + state.last_token_id = last_token_id + self._consume_token_sequences(state, last_token_id) + def _scan_prompt_state(self, prompt_slice: list[int]) -> tuple[bool, bool, int, Optional[int]]: started = False ended = False @@ -112,17 +235,6 @@ def update_state(self, share_inputs: dict) -> None: self._active_budgets = [] self._active_slots = [] - active_req_ids = [] - for req_id, stop_flag in zip(req_ids, stop_flags_list): - if stop_flag: - continue - if req_id: - active_req_ids.append(req_id) - - inactive_req_ids = set(self._states.keys()) - set(active_req_ids) - for req_id in inactive_req_ids: - self._states.pop(req_id, None) - candidate_slots = [] candidate_req_ids = [] candidate_args = [] @@ -133,8 +245,8 @@ def update_state(self, share_inputs: dict) -> None: if stop_flag or not req_id: continue - thinking_budget = logit_proc_args.get("thinking_budget") if logit_proc_args else None - if thinking_budget is None or not isinstance(thinking_budget, int) or thinking_budget < 0: + thinking_budget = self._resolve_thinking_budget(logit_proc_args, slot_id, share_inputs) + if thinking_budget is None: continue candidate_slots.append(slot_id) @@ -142,6 +254,10 @@ def update_state(self, share_inputs: dict) -> None: candidate_args.append(logit_proc_args) candidate_budgets.append(thinking_budget) + inactive_req_ids = set(self._states) - set(candidate_req_ids) + for req_id in inactive_req_ids: + del self._states[req_id] + if not candidate_slots: return @@ -176,6 +292,7 @@ def update_state(self, share_inputs: dict) -> None: thinking_budget = candidate_budgets[idx] state = self._states.setdefault(req_id, _ThinkingState()) + prompt_initialized = False if logit_proc_args: stop_sentence_token_ids = logit_proc_args.get("think_stop_sentence_token_ids") if isinstance(stop_sentence_token_ids, list) and all( @@ -202,6 +319,7 @@ def update_state(self, share_inputs: dict) -> None: if isinstance(prompt_last_token_id, int) and prompt_last_token_id >= 0: state.last_token_id = prompt_last_token_id state.prompt_checked = True + prompt_initialized = True current_step_idx = step_idx_by_slot.get(slot_id) state.current_step_idx = current_step_idx @@ -221,17 +339,23 @@ def update_state(self, share_inputs: dict) -> None: prompt_slice = list(prompt_slice) if prompt_ids is None: prompt_slice = [int(token_id) for token_id in prompt_slice if int(token_id) >= 0] - prompt_started, prompt_ended, prompt_tokens_after_start, prompt_last_token_id = ( - self._scan_prompt_state(prompt_slice) - ) - if prompt_started: - state.started = True - state.ended = prompt_ended - state.tokens_after_start = prompt_tokens_after_start - state.last_token_id = prompt_last_token_id + if self._sequence_mode: + self._scan_prompt_state_sequences(state, prompt_slice) if current_step_idx is not None and state.last_step_idx is None: state.last_step_idx = current_step_idx + else: + prompt_started, prompt_ended, prompt_tokens_after_start, prompt_last_token_id = ( + self._scan_prompt_state(prompt_slice) + ) + if prompt_started: + state.started = True + state.ended = prompt_ended + state.tokens_after_start = prompt_tokens_after_start + state.last_token_id = prompt_last_token_id + if current_step_idx is not None and state.last_step_idx is None: + state.last_step_idx = current_step_idx state.prompt_checked = True + prompt_initialized = True last_token_id = next_token_by_slot.get(slot_id) @@ -247,8 +371,10 @@ def update_state(self, share_inputs: dict) -> None: else: last_token_id = int(slot_pre_ids[-1].item()) - if last_token_id is not None and last_token_id >= 0: - if not state.started and last_token_id == self.think_start_token_id: + if last_token_id is not None and last_token_id >= 0 and not (self._sequence_mode and prompt_initialized): + if self._sequence_mode: + self._consume_decode_token_sequences(state, last_token_id, current_step_idx) + elif not state.started and last_token_id == self.think_start_token_id: state.started = True state.tokens_after_start = 0 state.last_token_id = last_token_id @@ -259,7 +385,7 @@ def update_state(self, share_inputs: dict) -> None: self._active_budgets.append(thinking_budget) self._active_slots.append(slot_id) continue - if current_step_idx is None: + elif current_step_idx is None: if last_token_id != state.last_token_id: state.last_token_id = last_token_id if state.started and not state.ended: @@ -301,6 +427,21 @@ def apply(self, logits: paddle.Tensor) -> paddle.Tensor: slot_id = self._active_slots[active_idx] stop_sentence_token_ids = state.stop_sentence_token_ids or [] stop_sentence_len = len(stop_sentence_token_ids) + if self._sequence_mode: + if not state.forcing: + if state.tokens_after_start < max(budget - stop_sentence_len, 0): + continue + state.forcing = True + state.forcing_pos = 0 + forced_sequence = stop_sentence_token_ids + self.think_forced_end_ids + force_token_id = forced_sequence[state.forcing_pos] + if force_token_id >= logits.shape[1]: + raise ValueError( + f"Forced thinking-end token {force_token_id} exceeds vocabulary size {logits.shape[1]}" + ) + logits[slot_id, :] = -float("inf") + logits[slot_id, force_token_id] = 0.0 + continue if stop_sentence_len > 0: budget_threshold = budget - stop_sentence_len if budget_threshold < 0: diff --git a/fastdeploy/model_executor/models/minicpm41/__init__.py b/fastdeploy/model_executor/models/minicpm41/__init__.py new file mode 100644 index 00000000000..7e917571e54 --- /dev/null +++ b/fastdeploy/model_executor/models/minicpm41/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .minicpm41 import MiniCPM41ForCausalLM + +__all__ = ["MiniCPM41ForCausalLM"] diff --git a/fastdeploy/model_executor/models/minicpm41/config_minicpm41.py b/fastdeploy/model_executor/models/minicpm41/config_minicpm41.py new file mode 100644 index 00000000000..9975b3fa559 --- /dev/null +++ b/fastdeploy/model_executor/models/minicpm41/config_minicpm41.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +SUPPORTED_QUANTIZATIONS = ("wint4", "wint8") diff --git a/fastdeploy/model_executor/models/minicpm41/hybrid_reasoning.py b/fastdeploy/model_executor/models/minicpm41/hybrid_reasoning.py new file mode 100644 index 00000000000..255ed60321b --- /dev/null +++ b/fastdeploy/model_executor/models/minicpm41/hybrid_reasoning.py @@ -0,0 +1,120 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from fastdeploy.model_executor.logits_processor.thinking_budget import ( + ThinkingBudgetLogitsProcessor, +) + +_DEFAULT_REASONING_TOKENS = "thinking" +_DEFAULT_MAX_THINKING_LENGTH = 512 + + +def _encode_text(tokenizer, text: str) -> list[int]: + token_ids = tokenizer.encode(text, add_special_tokens=False) + if hasattr(token_ids, "input_ids"): + token_ids = token_ids.input_ids + elif isinstance(token_ids, dict): + token_ids = token_ids["input_ids"] + if hasattr(token_ids, "tolist"): + token_ids = token_ids.tolist() + if len(token_ids) == 1 and isinstance(token_ids[0], list): + token_ids = token_ids[0] + return [int(token_id) for token_id in token_ids] + + +def build_minicpm41_thinking_token_sequences(tokenizer) -> dict: + """Build MiniCPM4.1 marker variants from its real tokenizer.""" + + def continuation_ids(prefix: str, suffix: str) -> list[int]: + prefix_ids = _encode_text(tokenizer, prefix) + combined_ids = _encode_text(tokenizer, prefix + suffix) + if combined_ids[: len(prefix_ids)] != prefix_ids: + raise ValueError(f"MiniCPM4.1 tokenizer cannot isolate {suffix!r} after {prefix!r}") + continuation = combined_ids[len(prefix_ids) :] + if not continuation: + raise ValueError(f"MiniCPM4.1 tokenizer produced an empty marker for {suffix!r}") + return continuation + + def unique(sequences: list[list[int]]) -> list[list[int]]: + return [list(sequence) for sequence in dict.fromkeys(tuple(sequence) for sequence in sequences)] + + start_sequences = unique([_encode_text(tokenizer, ""), continuation_ids("x", "")]) + end_sequences = unique([_encode_text(tokenizer, ""), continuation_ids("x", "")]) + if any(not sequence for sequence in start_sequences + end_sequences): + raise ValueError("MiniCPM4.1 tokenizer produced an empty thinking marker") + return { + "start": start_sequences, + "end": end_sequences, + "forced_end": continuation_ids("x", "\n\n"), + } + + +def _config_value(fd_config, name, default): + if hasattr(fd_config, name): + return getattr(fd_config, name) + model_config = getattr(fd_config, "model_config", None) + if model_config is not None and hasattr(model_config, name): + return getattr(model_config, name) + return default + + +def _flat_scalar(value) -> int: + if hasattr(value, "numpy"): + return int(value.numpy().reshape(-1)[0]) + if hasattr(value, "tolist"): + result = value.tolist() + return int(result[0] if isinstance(result, list) else result) + return int(value) + + +class HybridReasoningMode(ThinkingBudgetLogitsProcessor): + """MiniCPM4.1 hybrid reasoning mode. + + Reuses the generic ThinkingBudgetLogitsProcessor state machine; the + MiniCPM4.1-specific part is only the budget source: an explicit + `thinking_budget` request argument wins, then the request's + `reasoning_max_tokens` (share_inputs["max_think_lens"]), and finally the + model default `max_thinking_length`. Requests with `enable_thinking=False` + get no budget and are never forced. + """ + + def __init__(self, fd_config): + super().__init__(fd_config) + if not self._enabled: + raise ValueError("MiniCPM4.1 requires valid single-token markers or think_token_sequences") + self.reasoning_tokens = _config_value(fd_config, "reasoning_tokens", _DEFAULT_REASONING_TOKENS) + max_thinking_length = _config_value(fd_config, "max_thinking_length", _DEFAULT_MAX_THINKING_LENGTH) + if ( + isinstance(max_thinking_length, bool) + or not isinstance(max_thinking_length, int) + or max_thinking_length <= 0 + ): + raise ValueError("MiniCPM4.1 max_thinking_length must be a positive integer") + self.max_thinking_length = max_thinking_length + + def _resolve_thinking_budget(self, logit_proc_args, slot_id, share_inputs): + enable_thinking = share_inputs["enable_thinking"] + if not bool(_flat_scalar(enable_thinking[slot_id])): + return None + budget = super()._resolve_thinking_budget(logit_proc_args, slot_id, share_inputs) + if budget is not None: + return budget + max_think_lens = share_inputs["max_think_lens"] + configured = _flat_scalar(max_think_lens[slot_id]) + if configured == 0 or configured < -1: + raise ValueError("MiniCPM4.1 reasoning_max_tokens must be a positive integer or -1") + if configured > 0: + return configured + return self.max_thinking_length diff --git a/fastdeploy/model_executor/models/minicpm41/minicpm41.py b/fastdeploy/model_executor/models/minicpm41/minicpm41.py new file mode 100644 index 00000000000..30c552e721e --- /dev/null +++ b/fastdeploy/model_executor/models/minicpm41/minicpm41.py @@ -0,0 +1,457 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import math +import re +from functools import partial +from typing import Dict + +import paddle +from paddle import nn +from paddleformers.transformers import PretrainedModel +from paddleformers.utils.log import logger + +from fastdeploy.config import FDConfig, ModelConfig +from fastdeploy.model_executor.forward_meta import ForwardMeta +from fastdeploy.model_executor.graph_optimization.decorator import ( + support_graph_optimization, +) +from fastdeploy.model_executor.layers.activation import SiluAndMul +from fastdeploy.model_executor.layers.attention.attention import Attention +from fastdeploy.model_executor.layers.embeddings import VocabParallelEmbedding +from fastdeploy.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from fastdeploy.model_executor.layers.lm_head import ParallelLMHead +from fastdeploy.model_executor.layers.normalization import RMSNorm +from fastdeploy.model_executor.models.minicpm41.config_minicpm41 import ( + SUPPORTED_QUANTIZATIONS, +) +from fastdeploy.model_executor.models.minicpm41.hybrid_reasoning import ( + HybridReasoningMode, + build_minicpm41_thinking_token_sequences, +) +from fastdeploy.model_executor.models.model_base import ( + ModelCategory, + ModelForCasualLM, + ModelRegistry, +) +from fastdeploy.model_executor.utils import ( + WeightsMapper, + default_weight_loader, + process_weights_after_loading, + process_weights_before_loading, +) + + +def minicpm41_config_value(model_config, name: str, default): + value = getattr(model_config, name, None) + if value is not None: + return value + + pretrained_config = getattr(model_config, "pretrained_config", None) + value = getattr(pretrained_config, name, None) + if value is not None: + return value + + return default + + +def minicpm41_embedding_scale(model_config) -> float: + return float(minicpm41_config_value(model_config, "scale_emb", 1.0)) + + +def minicpm41_residual_scale(model_config) -> float: + num_hidden_layers = int(minicpm41_config_value(model_config, "num_hidden_layers", 1)) + scale_depth = float(minicpm41_config_value(model_config, "scale_depth", 1.0)) + return scale_depth / math.sqrt(num_hidden_layers) + + +def minicpm41_lm_head_scale(model_config) -> float: + hidden_size = float(minicpm41_config_value(model_config, "hidden_size", 1.0)) + dim_model_base = float(minicpm41_config_value(model_config, "dim_model_base", hidden_size)) + return dim_model_base / hidden_size + + +class MiniCPM41MLP(nn.Layer): + """MiniCPM4.1 feed-forward network.""" + + def __init__(self, fd_config: FDConfig, layer_id: int, prefix: str = "") -> None: + super().__init__() + self.up_gate_proj = MergedColumnParallelLinear( + fd_config=fd_config, + prefix=f"{prefix}.up_gate_proj", + input_size=fd_config.model_config.hidden_size, + output_size=fd_config.model_config.intermediate_size * 2, + with_bias=False, + activation=fd_config.model_config.hidden_act, + ) + + self.down_proj = RowParallelLinear( + fd_config=fd_config, + prefix=f"{prefix}.down_proj", + input_size=fd_config.model_config.intermediate_size, + output_size=fd_config.model_config.hidden_size, + with_bias=False, + ) + + self.act_fn = SiluAndMul( + fd_config=fd_config, + bias=getattr(self.up_gate_proj, "bias", None), + act_method=fd_config.model_config.hidden_act, + ) + + def load_state_dict(self, state_dict): + self.up_gate_proj.load_state_dict(state_dict) + self.down_proj.load_state_dict(state_dict) + + def forward(self, x: paddle.Tensor, forward_meta: ForwardMeta): + gate_up_out = self.up_gate_proj(x) + act_out = self.act_fn(gate_up_out) + return self.down_proj(act_out) + + +class MiniCPM41Attention(nn.Layer): + """MiniCPM4.1 grouped-query self attention.""" + + def __init__(self, fd_config: FDConfig, layer_id: int, prefix: str = "") -> None: + super().__init__() + self.qkv_proj = QKVParallelLinear(fd_config=fd_config, prefix=f"{prefix}.qkv_proj", with_bias=False) + + self.o_proj = RowParallelLinear( + fd_config=fd_config, + prefix=f"{prefix}.o_proj", + input_size=fd_config.model_config.hidden_size, + output_size=fd_config.model_config.hidden_size, + layer_id=layer_id, + ) + + self.attn = Attention( + fd_config=fd_config, + layer_id=layer_id, + prefix=prefix, + use_neox_rotary_style=True, + ) + + def load_state_dict(self, state_dict): + self.qkv_proj.load_state_dict(state_dict) + self.o_proj.load_state_dict(state_dict) + self.attn.load_state_dict(state_dict) + + def forward( + self, + forward_meta: ForwardMeta, + hidden_states: paddle.Tensor, + ): + qkv_out = self.qkv_proj(hidden_states) + attention_out = self.attn(qkv=qkv_out, forward_meta=forward_meta) + return self.o_proj(attention_out) + + +class MiniCPM41DecoderLayer(nn.Layer): + """MiniCPM4.1 decoder block with muP residual scaling.""" + + def __init__(self, fd_config: FDConfig, prefix: str = "") -> None: + super().__init__() + layer_id = int(prefix.split(sep=".")[-1]) + self.residual_scale = minicpm41_residual_scale(fd_config.model_config) + + self.self_attn = MiniCPM41Attention( + fd_config=fd_config, + layer_id=layer_id, + prefix=f"{prefix}.self_attn", + ) + + self.mlp = MiniCPM41MLP( + fd_config=fd_config, + layer_id=layer_id, + prefix=f"{prefix}.mlp", + ) + + self.input_layernorm = RMSNorm( + fd_config, + hidden_size=fd_config.model_config.hidden_size, + eps=fd_config.model_config.rms_norm_eps, + prefix=f"{prefix}.input_layernorm", + ) + + self.post_attention_layernorm = RMSNorm( + fd_config, + hidden_size=fd_config.model_config.hidden_size, + eps=fd_config.model_config.rms_norm_eps, + prefix=f"{prefix}.post_attention_layernorm", + layer_id=layer_id, + ) + + def load_state_dict(self, state_dict): + self.self_attn.load_state_dict(state_dict) + self.mlp.load_state_dict(state_dict) + self.input_layernorm.load_state_dict(state_dict) + self.post_attention_layernorm.load_state_dict(state_dict) + + def forward( + self, + forward_meta: ForwardMeta, + hidden_states: paddle.Tensor, + residual: paddle.Tensor = None, + ): + hidden_states, residual = self.input_layernorm( + hidden_states, residual_input=residual, forward_meta=forward_meta + ) + hidden_states = self.self_attn(forward_meta=forward_meta, hidden_states=hidden_states) + hidden_states = hidden_states * self.residual_scale + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states, forward_meta) + hidden_states = hidden_states * self.residual_scale + + return hidden_states, residual + + +@support_graph_optimization +class MiniCPM41Model(nn.Layer): + """MiniCPM4.1 decoder-only transformer.""" + + def __init__(self, fd_config: FDConfig = None): + super().__init__() + + self.num_layers = fd_config.model_config.num_hidden_layers + self.embedding_scale = minicpm41_embedding_scale(fd_config.model_config) + fd_config.model_config.pretrained_config.prefix_name = "minicpm41" + + self.embed_tokens = VocabParallelEmbedding( + fd_config=fd_config, + num_embeddings=fd_config.model_config.vocab_size, + embedding_dim=fd_config.model_config.hidden_size, + params_dtype=paddle.get_default_dtype, + prefix=f"{fd_config.model_config.pretrained_config.prefix_name}.embed_tokens", + ) + + self.layers = nn.LayerList( + [ + MiniCPM41DecoderLayer( + fd_config=fd_config, + prefix=f"{fd_config.model_config.pretrained_config.prefix_name}.layers.{i}", + ) + for i in range(self.num_layers) + ] + ) + + self.norm = RMSNorm( + fd_config, + hidden_size=fd_config.model_config.hidden_size, + eps=fd_config.model_config.rms_norm_eps, + prefix=f"{fd_config.model_config.pretrained_config.prefix_name}.norm", + ) + + def load_state_dict(self, state_dict): + self.embed_tokens.load_state_dict(state_dict) + self.norm.load_state_dict(state_dict) + for i in range(self.num_layers): + logger.info(f"Start load layer {i}") + self.layers[i].load_state_dict(state_dict) + + def forward( + self, + ids_remove_padding: paddle.Tensor, + forward_meta: ForwardMeta, + ): + hidden_states = self.embed_tokens(ids_remove_padding=ids_remove_padding, forward_meta=forward_meta) + hidden_states = hidden_states * self.embedding_scale + + residual = None + for i in range(self.num_layers): + hidden_states, residual = self.layers[i](forward_meta, hidden_states, residual) + + return self.norm(hidden_states, residual)[0] + + +@ModelRegistry.register_model_class( + architecture="MiniCPMForCausalLM", + module_name="minicpm41.minicpm41", + category=ModelCategory.TEXT_GENERATION | ModelCategory.REASONING, +) +class MiniCPM41ForCausalLM(ModelForCasualLM): + """MiniCPM4.1-8B FastDeploy entry point.""" + + supported_quantizations = SUPPORTED_QUANTIZATIONS + + def __init__(self, fd_config: FDConfig): + super().__init__(fd_config) + self.fd_config = fd_config + self.hybrid_reasoning = HybridReasoningMode(fd_config) + self.minicpm41 = MiniCPM41Model(fd_config=fd_config) + + self.ori_vocab_size = fd_config.model_config.ori_vocab_size + self.tie_word_embeddings = fd_config.model_config.tie_word_embeddings + self.lm_head_scale = minicpm41_lm_head_scale(fd_config.model_config) + self.lm_head = ParallelLMHead( + fd_config=fd_config, + embedding_dim=fd_config.model_config.hidden_size, + num_embeddings=fd_config.model_config.vocab_size, + prefix="lm_head", + ) + + self.process_weights_before_loading_fn = process_weights_before_loading( + mapper=( + WeightsMapper(orig_to_new_prefix={"model.": "minicpm41."}) + if self.fd_config.model_config.model_format == "torch" + else None + ), + ) + + @classmethod + def name(self): + return "MiniCPMForCausalLM" + + @staticmethod + def build_thinking_token_sequences(tokenizer): + """Expose MiniCPM4.1 tokenizer markers through the engine model hook.""" + return build_minicpm41_thinking_token_sequences(tokenizer) + + @paddle.no_grad() + def load_weights(self, weights_iterator) -> None: + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("up_gate_proj", "gate_proj", "gate"), + ("up_gate_proj", "up_proj", "up"), + ("embed_tokens.embeddings", "embed_tokens", None), + ("lm_head.linear", "lm_head", None), + ] + + params_dict = dict(self.named_parameters()) + process_weights_after_loading_fn = process_weights_after_loading(dict(self.named_sublayers()), self.fd_config) + + for loaded_weight_name, loaded_weight in weights_iterator: + logger.debug(f"Loading weight: {loaded_weight_name}") + loaded_weight_name = ( + self.process_weights_before_loading_fn(loaded_weight_name) + if getattr(self, "process_weights_before_loading_fn", None) + else loaded_weight_name + ) + if loaded_weight_name is None: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in loaded_weight_name: + continue + model_param_name = loaded_weight_name.replace(weight_name, param_name) + if model_param_name not in params_dict: + continue + param = params_dict[model_param_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader(self.fd_config)) + weight_loader(param, loaded_weight, shard_id) + break + else: + model_param_name = loaded_weight_name + if model_param_name not in params_dict: + continue + param = params_dict[model_param_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader(self.fd_config)) + weight_loader(param, loaded_weight) + + model_sublayer_name = re.sub(r"\.(weight|weight_scale)$", "", model_param_name) + process_weights_after_loading_fn(model_sublayer_name, param) + + if getattr(self, "tie_word_embeddings", False): + self.lm_head.linear.weight.set_value( + self.minicpm41.embed_tokens.embeddings.weight.transpose([1, 0]).astype( + self.lm_head.linear.weight.dtype + ) + ) + + @paddle.no_grad() + def set_state_dict(self, state_dict): + self.minicpm41.load_state_dict(state_dict) + self.lm_head.load_state_dict(state_dict) + + def compute_logits(self, hidden_states: paddle.Tensor, forward_meta: ForwardMeta = None): + hidden_states = hidden_states * self.lm_head_scale + logits = self.lm_head(hidden_states) + logits = logits.astype(paddle.float32) + logits[:, self.ori_vocab_size :] = -float("inf") + return logits + + def get_logits_processors(self): + """Return model-owned logits processors for the runner sampling path.""" + return [self.hybrid_reasoning] + + def forward( + self, + inputs: Dict, + forward_meta: ForwardMeta, + ): + ids_remove_padding = inputs["ids_remove_padding"] + return self.minicpm41(ids_remove_padding=ids_remove_padding, forward_meta=forward_meta) + + def clear_graph_opt_backend(self): + """Clear graph optimization backend, the captured cuda graph will be cleaned""" + self.minicpm41.clear_graph_opt_backend(fd_config=self.fd_config) + + +class MiniCPM41PretrainedModel(PretrainedModel): + """MiniCPM4.1 tensor-parallel conversion metadata.""" + + config_class = FDConfig + + def _init_weight(self, layer): + return None + + @classmethod + def arch_name(self): + return "MiniCPMForCausalLM" + + @classmethod + def _get_tensor_parallel_mappings(cls, config: ModelConfig, is_split=True): + from paddleformers.transformers.conversion_utils import split_or_merge_func + + fn = split_or_merge_func( + is_split=is_split, + tensor_model_parallel_size=config.tensor_model_parallel_size, + tensor_parallel_rank=config.tensor_parallel_rank, + num_attention_heads=config.num_attention_heads, + ) + + def get_tensor_parallel_split_mappings(num_layers): + final_actions = {} + + base_actions = { + "lm_head.weight": partial(fn, is_column=True), + "embed_tokens.weight": partial(fn, is_column=False), + "layers.0.self_attn.o_proj.weight": partial(fn, is_column=False), + "layers.0.mlp.down_proj.weight": partial(fn, is_column=False), + "layers.0.self_attn.q_proj.weight": partial(fn, is_column=True), + "layers.0.mlp.gate_proj.weight": partial(fn, is_column=True), + "layers.0.mlp.up_proj.weight": partial(fn, is_column=True), + } + + if config.num_key_value_heads % config.tensor_model_parallel_size == 0: + base_actions["layers.0.self_attn.k_proj.weight"] = partial(fn, is_column=True) + base_actions["layers.0.self_attn.v_proj.weight"] = partial(fn, is_column=True) + + for key, action in base_actions.items(): + if "layers.0." in key: + for i in range(num_layers): + final_actions[key.replace("layers.0.", f"layers.{i}.")] = action + final_actions[key] = action + + return final_actions + + return get_tensor_parallel_split_mappings(config.num_hidden_layers) diff --git a/fastdeploy/platforms/base.py b/fastdeploy/platforms/base.py index b2eceb0aeb8..fb71ec96931 100644 --- a/fastdeploy/platforms/base.py +++ b/fastdeploy/platforms/base.py @@ -31,6 +31,7 @@ class _Backend(enum.Enum): HPU_ATTN = enum.auto() FLASH_MASK_ATTN = enum.auto() DECODE_UNIFIED_ATTN = enum.auto() + INFLLMV2_ATTN = enum.auto() class Platform: diff --git a/fastdeploy/platforms/cuda.py b/fastdeploy/platforms/cuda.py index e9a3cb61574..88e981029a2 100644 --- a/fastdeploy/platforms/cuda.py +++ b/fastdeploy/platforms/cuda.py @@ -76,8 +76,12 @@ def get_attention_backend_cls(cls, selected_backend: _Backend): elif selected_backend == _Backend.DECODE_UNIFIED_ATTN: logger.info("Using DECODE UNIFIED ATTN backend.") return "fastdeploy.model_executor.layers.attention.DecodeUnifiedAttentionBackend" + elif selected_backend == _Backend.INFLLMV2_ATTN: + logger.info("Using INFLLMV2 ATTN backend.") + return "fastdeploy.model_executor.layers.attention.InfLLMV2AttentionBackend" else: raise ValueError( "Invalid attention backend you specified.\n" - "Now only support [NATIVE_ATTN, MLA_ATTN, APPEND_ATTN, DECODE_UNIFIED_ATTN, FLASH_ATTN] in cuda place." + "Now only support [NATIVE_ATTN, MLA_ATTN, APPEND_ATTN, DECODE_UNIFIED_ATTN, " + "INFLLMV2_ATTN, FLASH_ATTN] in cuda place." ) diff --git a/fastdeploy/worker/model_runner_base.py b/fastdeploy/worker/model_runner_base.py index a050391d010..a8b4a14b514 100644 --- a/fastdeploy/worker/model_runner_base.py +++ b/fastdeploy/worker/model_runner_base.py @@ -78,6 +78,19 @@ def get_model(self) -> nn.Layer: """ raise NotImplementedError + def register_model_logits_processors(self) -> None: + """Register optional logits processors owned by the loaded model.""" + provider = getattr(self.get_model(), "get_logits_processors", None) + if provider is None: + return + processors = provider() + if not isinstance(processors, (list, tuple)): + raise TypeError("get_logits_processors() must return a list or tuple") + registered = self.share_inputs["logits_processors"] + for processor in processors: + if all(processor is not current for current in registered): + registered.append(processor) + @abstractmethod def execute_model( self, diff --git a/fastdeploy/worker/worker_process.py b/fastdeploy/worker/worker_process.py index 9b5cb67f731..a988526ce42 100644 --- a/fastdeploy/worker/worker_process.py +++ b/fastdeploy/worker/worker_process.py @@ -716,6 +716,7 @@ def load_model(self) -> None: """Load weights and create model""" self.worker.load_model() + self.worker.model_runner.register_model_logits_processors() loaded_model_signal_data = np.zeros(shape=[1], dtype=np.int32) self.loaded_model_signal = IPCSignal( name="loaded_model_signal", @@ -886,6 +887,7 @@ def parse_args(): parser.add_argument("--ori_vocab_size", type=int, default=None) parser.add_argument("--think_start_id", type=int, default=-1) parser.add_argument("--think_end_id", type=int, default=-1) + parser.add_argument("--think_token_sequences", type=json.loads, default=None) parser.add_argument("--image_patch_id", type=int, default=-1) parser.add_argument("--line_break_id", type=int, default=-1) parser.add_argument("--think_truncate_prompt_ids", type=json.loads, default=[]) diff --git a/scripts/evaluate_minicpm41_quant_quality.py b/scripts/evaluate_minicpm41_quant_quality.py new file mode 100644 index 00000000000..4a782e46bef --- /dev/null +++ b/scripts/evaluate_minicpm41_quant_quality.py @@ -0,0 +1,224 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run the deterministic MiniCPM4.1 quantization quality smoke suite.""" + +from __future__ import annotations + +import argparse +import json +import time +import unicodedata +from pathlib import Path +from urllib.error import HTTPError, URLError +from urllib.request import ProxyHandler, Request, build_opener + +QUALITY_CASES = ( + { + "name": "integer_multiplication", + "prompt": "请只输出阿拉伯数字，不要解释：17乘以23等于多少？", + "exact_answers": ("391",), + }, + { + "name": "world_capital", + "prompt": "请只输出城市名，不要解释：法国的首都是什么？", + "exact_answers": ("巴黎",), + }, + { + "name": "water_formula", + "prompt": "请只输出化学式，不要解释：水的化学式是什么？", + "exact_answers": ("H2O",), + }, + { + "name": "square_is_rectangle", + "prompt": "请只回答“是”或“否”，不要解释：所有正方形都是矩形吗？", + "exact_answers": ("是",), + }, + { + "name": "number_sequence", + "prompt": "请只输出下一个数字，不要解释：数列2、4、8、16的下一项是什么？", + "exact_answers": ("32",), + }, + { + "name": "rayleigh_scattering", + "prompt": "请用一句话解释为什么天空是蓝色的。", + "required_terms": ("蓝", "散射"), + }, + { + "name": "integer_addition", + "prompt": "请只输出阿拉伯数字，不要解释：125加376等于多少？", + "exact_answers": ("501",), + }, + { + "name": "square_root", + "prompt": "请只输出阿拉伯数字，不要解释：81的算术平方根是多少？", + "exact_answers": ("9",), + }, + { + "name": "largest_planet", + "prompt": "请只输出行星名，不要解释：太阳系中体积最大的行星是什么？", + "exact_answers": ("木星",), + }, + { + "name": "binary_conversion", + "prompt": "请只输出阿拉伯数字，不要解释：二进制1010对应的十进制数是多少？", + "exact_answers": ("10",), + }, + { + "name": "literature_author", + "prompt": "请只输出人名，不要解释：《红楼梦》的作者通常认为是谁？", + "exact_answers": ("曹雪芹",), + }, + { + "name": "basic_syllogism", + "prompt": "请只回答“是”或“否”，不要解释：所有猫都是动物，小花是一只猫，因此小花是动物，对吗？", + "exact_answers": ("是",), + }, +) + + +def generation_config(max_tokens: int) -> dict: + return { + "max_tokens": max_tokens, + "temperature": 0, + "top_p": 1, + "stream": False, + "chat_template_kwargs": {"enable_thinking": False}, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="Served model name or local checkpoint path.") + parser.add_argument("--label", required=True, help="Stable report label, for example bf16 or wint4.") + parser.add_argument("--base-url", default="http://127.0.0.1:8180", help="FastDeploy API base URL.") + parser.add_argument("--output", type=Path, required=True, help="Path for the JSON report.") + parser.add_argument("--max-tokens", type=int, default=96) + parser.add_argument("--timeout", type=float, default=180.0) + parser.add_argument("--min-passed", type=int, default=10) + return parser.parse_args() + + +def normalize_short_answer(output: str) -> str: + normalized = unicodedata.normalize("NFKC", output).strip() + return normalized.rstrip("。.!！").strip() + + +def score_output(case: dict, output: str) -> tuple[bool, str]: + normalized = normalize_short_answer(output) + exact_answers = case.get("exact_answers") + if exact_answers is not None: + expected = tuple(normalize_short_answer(answer) for answer in exact_answers) + passed = normalized in expected + return passed, f"expected one of {expected}, got {normalized!r}" + + required_terms = tuple(case["required_terms"]) + missing = [term for term in required_terms if term not in normalized] + return not missing, "missing required terms: " + ", ".join(missing) if missing else "all required terms found" + + +def send_chat(args: argparse.Namespace, prompt: str) -> tuple[dict, float]: + payload = { + "model": args.model, + "messages": [{"role": "user", "content": prompt}], + **generation_config(args.max_tokens), + } + request = Request( + f"{args.base_url.rstrip('/')}/v1/chat/completions", + data=json.dumps(payload, ensure_ascii=False).encode("utf-8"), + headers={"Content-Type": "application/json"}, + method="POST", + ) + start = time.perf_counter() + try: + with build_opener(ProxyHandler({})).open(request, timeout=args.timeout) as response: + result = json.loads(response.read().decode("utf-8")) + except HTTPError as error: + body = error.read().decode("utf-8", errors="replace") + raise RuntimeError(f"server returned HTTP {error.code}: {body}") from error + except URLError as error: + raise RuntimeError(f"cannot reach {args.base_url}: {error.reason}") from error + return result, time.perf_counter() - start + + +def extract_output(response: dict) -> str: + try: + output = response["choices"][0]["message"]["content"] + except (KeyError, IndexError, TypeError) as error: + raise ValueError(f"unexpected chat completion response: {response}") from error + if not isinstance(output, str): + raise ValueError(f"chat completion content must be a string, got {type(output).__name__}") + return output + + +def run_suite(args: argparse.Namespace) -> dict: + if not 1 <= args.min_passed <= len(QUALITY_CASES): + raise ValueError(f"--min-passed must be between 1 and {len(QUALITY_CASES)}") + + results = [] + for case in QUALITY_CASES: + response, elapsed = send_chat(args, case["prompt"]) + output = extract_output(response) + passed, detail = score_output(case, output) + results.append( + { + "name": case["name"], + "prompt": case["prompt"], + "output": output, + "passed": passed, + "score_detail": detail, + "latency_seconds": elapsed, + "usage": response.get("usage", {}), + } + ) + + passed = sum(result["passed"] for result in results) + return { + "label": args.label, + "model": args.model, + "generation_config": generation_config(args.max_tokens), + "threshold": {"min_passed": args.min_passed, "total": len(QUALITY_CASES)}, + "summary": { + "passed": passed, + "total": len(QUALITY_CASES), + "pass_rate": passed / len(QUALITY_CASES), + "accepted": passed >= args.min_passed, + }, + "cases": results, + } + + +def main() -> int: + args = parse_args() + try: + report = run_suite(args) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + except (OSError, RuntimeError, ValueError, json.JSONDecodeError) as error: + print(f"error: {error}") + return 1 + + summary = report["summary"] + for case in report["cases"]: + state = "PASS" if case["passed"] else "FAIL" + print(f"[{state}] {case['name']}: {case['output']}") + print( + f"{args.label}: {summary['passed']}/{summary['total']} passed; " + f"accepted={summary['accepted']}; report={args.output}" + ) + return 0 if summary["accepted"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_minicpm41_wint_server.sh b/scripts/run_minicpm41_wint_server.sh new file mode 100755 index 00000000000..1a11a1f7958 --- /dev/null +++ b/scripts/run_minicpm41_wint_server.sh @@ -0,0 +1,115 @@ +#!/usr/bin/env bash + +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" + +QUANTIZATION="${QUANTIZATION:-}" +if [[ -z "${QUANTIZATION}" && $# -gt 0 ]]; then + QUANTIZATION="$1" + shift +fi +case "${QUANTIZATION}" in + wint4 | wint8) ;; + *) + printf 'Quantization must be wint4 or wint8, got: %s\n' "${QUANTIZATION:-}" >&2 + printf 'Usage: %s {wint4|wint8} /path/to/MiniCPM4.1-8B [server options]\n' "$0" >&2 + exit 2 + ;; +esac + +MODEL_PATH="${MODEL_PATH:-}" +if [[ -z "${MODEL_PATH}" && $# -gt 0 && "${1}" != --* ]]; then + MODEL_PATH="$1" + shift +fi +if [[ -z "${MODEL_PATH}" ]]; then + printf 'Model path is required.\n' >&2 + printf 'Usage: %s {wint4|wint8} /path/to/MiniCPM4.1-8B [server options]\n' "$0" >&2 + exit 2 +fi + +PYTHON_BIN="${PYTHON_BIN:-${REPO_ROOT}/.venv/bin/python}" +PYTHON_ENV_BIN="$(cd "$(dirname "${PYTHON_BIN}")" && pwd)" +PYTHON_ENV_ROOT="$(cd "${PYTHON_ENV_BIN}/.." && pwd)" +CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}" +GPU_IDS="${CUDA_VISIBLE_DEVICES:-0}" +IFS=',' read -r -a GPU_LIST <<< "${GPU_IDS}" +TP_SIZE="${TP_SIZE:-${#GPU_LIST[@]}}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-8192}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-1}" +API_PORT="${FD_API_PORT:-8180}" +METRICS_PORT="${FD_METRICS_PORT:-8181}" +ENGINE_QUEUE_PORT="${FD_ENGINE_QUEUE_PORT:-8182}" +CACHE_QUEUE_PORT="${FD_CACHE_QUEUE_PORT:-8183}" +ATTENTION_BACKEND="${FD_ATTENTION_BACKEND:-FLASH_ATTN}" +MODEL_SOURCE="${FD_MODEL_SOURCE:-HUGGINGFACE}" +SERVER_WORKERS="${FD_SERVER_WORKERS:-1}" +SERVED_MODEL_NAME="${FD_SERVED_MODEL_NAME:-MiniCPM4.1-8B}" +MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-128}" + +if [[ ! -x "${PYTHON_BIN}" ]]; then + printf 'Python executable not found: %s\n' "${PYTHON_BIN}" >&2 + exit 2 +fi +if [[ ! -d "${MODEL_PATH}" ]]; then + printf 'Model directory not found: %s\n' "${MODEL_PATH}" >&2 + exit 2 +fi + +if [[ -x "${CUDA_HOME}/bin/nvcc" ]]; then + export PATH="${CUDA_HOME}/bin:${PYTHON_ENV_BIN}:${PATH}" +fi +PYTHON_NVIDIA_LIBRARY_PATH="$( + find "${PYTHON_ENV_ROOT}/lib" \ + -type d -path '*/site-packages/nvidia/*/lib' -print | sort | paste -sd: - +)" +if [[ -n "${PYTHON_NVIDIA_LIBRARY_PATH}" ]]; then + export LD_LIBRARY_PATH="${PYTHON_NVIDIA_LIBRARY_PATH}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" +fi + +PYTHONPATH_VALUE="${REPO_ROOT}" +if [[ -n "${PYTHONPATH:-}" ]]; then + PYTHONPATH_VALUE="${PYTHONPATH_VALUE}:${PYTHONPATH}" +fi + +printf 'Starting MiniCPM4.1 %s online-quantized server\n' "${QUANTIZATION}" +printf ' model: %s\n' "${MODEL_PATH}" +printf ' GPUs: %s (TP=%s)\n' "${GPU_IDS}" "${TP_SIZE}" +printf ' API: http://127.0.0.1:%s/v1/chat/completions\n' "${API_PORT}" + +exec env \ + PYTHONPATH="${PYTHONPATH_VALUE}" \ + CUDA_VISIBLE_DEVICES="${GPU_IDS}" \ + FD_ATTENTION_BACKEND="${ATTENTION_BACKEND}" \ + FD_MODEL_SOURCE="${MODEL_SOURCE}" \ + "${PYTHON_BIN}" -u -m fastdeploy.entrypoints.openai.api_server \ + --model "${MODEL_PATH}" \ + --port "${API_PORT}" \ + --metrics-port "${METRICS_PORT}" \ + --engine-worker-queue-port "${ENGINE_QUEUE_PORT}" \ + --cache-queue-port "${CACHE_QUEUE_PORT}" \ + --tensor-parallel-size "${TP_SIZE}" \ + --max-model-len "${MAX_MODEL_LEN}" \ + --max-num-seqs "${MAX_NUM_SEQS}" \ + --max-num-batched-tokens "${MAX_NUM_BATCHED_TOKENS}" \ + --workers "${SERVER_WORKERS}" \ + --served-model-name "${SERVED_MODEL_NAME}" \ + --no-enable-prefix-caching \ + --quantization "${QUANTIZATION}" \ + "$@" diff --git a/tests/benchmarks/test_minicpm41.sh b/tests/benchmarks/test_minicpm41.sh new file mode 100755 index 00000000000..41e39398276 --- /dev/null +++ b/tests/benchmarks/test_minicpm41.sh @@ -0,0 +1,157 @@ +#!/usr/bin/env bash + +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" +cd "${REPO_ROOT}" + +TARGET="${1:-all}" +PYTHON_BIN="${PYTHON_BIN:-${REPO_ROOT}/.venv/bin/python}" +PYTHON_ENV_BIN="$(cd "$(dirname "${PYTHON_BIN}")" && pwd)" +PYTHON_ENV_ROOT="$(cd "${PYTHON_ENV_BIN}/.." && pwd)" +CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}" +GPU_IDS="${CUDA_VISIBLE_DEVICES:-0}" +MODEL_PATH="${MINICPM41_MODEL_PATH:-${MODEL_PATH:-}}" + +case "${TARGET}" in + all | build | unit | operators | e2e) ;; + *) + printf 'Usage: %s [all|build|unit|operators|e2e]\n' "$0" >&2 + exit 2 + ;; +esac + +if [[ ! -x "${PYTHON_BIN}" ]]; then + printf 'Python executable not found: %s\n' "${PYTHON_BIN}" >&2 + exit 2 +fi + +export PATH="${CUDA_HOME}/bin:${PYTHON_ENV_BIN}:${PATH}" +export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" +PYTHON_NVIDIA_LIBRARY_PATH="$( + find "${PYTHON_ENV_ROOT}/lib" \ + -type d -path '*/site-packages/nvidia/*/lib' -print | sort | paste -sd: - +)" +if [[ -n "${PYTHON_NVIDIA_LIBRARY_PATH}" ]]; then + export LD_LIBRARY_PATH="${PYTHON_NVIDIA_LIBRARY_PATH}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" +fi + +require_model() { + if [[ -z "${MODEL_PATH}" || ! -d "${MODEL_PATH}" ]]; then + printf 'Set MODEL_PATH to the local MiniCPM4.1-8B directory.\n' >&2 + return 2 + fi +} + +verify_build() { + local library + local symbol + library="$(find fastdeploy/model_executor/ops/gpu -type f -name 'fastdeploy_ops*.so' -print -quit)" + if [[ -z "${library}" ]]; then + printf 'fastdeploy_ops shared library was not generated.\n' >&2 + return 1 + fi + for symbol in infllmv2_update_compressed_k infllmv2_select_blocks infllmv2_attention_forward; do + strings "${library}" | grep "${symbol}" >/dev/null + done + printf 'PASS build: %s\n' "${library}" +} + +run_build() { + local cudnn_header + local nvidia_ml_stub="" + local candidate + cudnn_header="$( + find "${PYTHON_ENV_ROOT}/lib" \ + -type f -path '*/site-packages/nvidia/cudnn/include/cudnn.h' -print -quit + )" + if [[ -n "${cudnn_header}" ]]; then + export CPATH="$(dirname "${cudnn_header}")${CPATH:+:${CPATH}}" + fi + for candidate in \ + "${CUDA_HOME}/targets/x86_64-linux/lib/stubs" \ + /usr/lib/x86_64-linux-gnu/stubs; do + if [[ -f "${candidate}/libnvidia-ml.so" ]]; then + nvidia_ml_stub="${candidate}" + break + fi + done + if [[ -n "${nvidia_ml_stub}" ]]; then + export LIBRARY_PATH="${nvidia_ml_stub}:${PYTHON_NVIDIA_LIBRARY_PATH}${LIBRARY_PATH:+:${LIBRARY_PATH}}" + fi + MAX_JOBS="${MAX_JOBS:-4}" \ + bash build.sh 0 "${PYTHON_BIN}" false "${FD_BUILDING_ARCS:-[86]}" + verify_build +} + +run_unit() { + "${PYTHON_BIN}" -m pytest -q \ + tests/model_executor/test_minicpm41.py \ + tests/model_executor/test_thinking_budget.py \ + tests/model_executor/test_infllmv2_attention_backend.py \ + tests/quantization/test_minicpm41_int_quant.py \ + tests/quantization/test_minicpm41_quality_eval.py +} + +run_operators() { + CUDA_VISIBLE_DEVICES="${GPU_IDS}" "${PYTHON_BIN}" -m pytest -q \ + tests/operators/test_infllmv2_attention_forward.py +} + +run_e2e_nodes() { + "${PYTHON_BIN}" -m pytest -q \ + tests/e2e/test_minicpm41_serving.py::test_minicpm41_openai_chat_completion_e2e \ + tests/e2e/test_minicpm41_serving.py::test_minicpm41_chat_completion_with_history_e2e \ + tests/e2e/test_minicpm41_serving.py::test_minicpm41_forces_multitoken_thinking_end_e2e \ + tests/e2e/test_minicpm41_serving.py::test_minicpm41_disable_thinking_does_not_emit_think_block_e2e \ + tests/e2e/test_minicpm41_serving.py::test_minicpm41_mixed_thinking_modes_e2e +} + +run_e2e() { + require_model + export MINICPM41_MODEL_PATH="${MODEL_PATH}" + export MINICPM41_E2E_MAX_MODEL_LEN="${MINICPM41_E2E_MAX_MODEL_LEN:-1024}" + export MINICPM41_E2E_MAX_NUM_SEQS="${MINICPM41_E2E_MAX_NUM_SEQS:-2}" + export MINICPM41_E2E_EXTRA_ARGS="${MINICPM41_E2E_EXTRA_ARGS:---workers 1 --max-num-batched-tokens 256 --num-gpu-blocks-override 32 --no-enable-prefix-caching}" + export CUDA_VISIBLE_DEVICES="${GPU_IDS}" + + export FD_ATTENTION_BACKEND=FLASH_ATTN + unset MINICPM41_E2E_QUANTIZATION + run_e2e_nodes + export MINICPM41_E2E_QUANTIZATION=wint4 + run_e2e_nodes + export MINICPM41_E2E_QUANTIZATION=wint8 + run_e2e_nodes + unset MINICPM41_E2E_QUANTIZATION + export FD_ATTENTION_BACKEND=INFLLMV2_ATTN + run_e2e_nodes + printf 'PASS e2e: BF16, WINT4, WINT8, and InfLLM-V2\n' +} + +if [[ "${TARGET}" == "all" || "${TARGET}" == "build" ]]; then + run_build +fi +if [[ "${TARGET}" == "all" || "${TARGET}" == "unit" ]]; then + run_unit +fi +if [[ "${TARGET}" == "all" || "${TARGET}" == "operators" ]]; then + run_operators +fi +if [[ "${TARGET}" == "all" || "${TARGET}" == "e2e" ]]; then + run_e2e +fi diff --git a/tests/e2e/test_minicpm41_serving.py b/tests/e2e/test_minicpm41_serving.py new file mode 100644 index 00000000000..c621a828567 --- /dev/null +++ b/tests/e2e/test_minicpm41_serving.py @@ -0,0 +1,522 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +import os +import shlex +import signal +import socket +import subprocess +import sys +import time +import urllib.error +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, replace +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +LOCAL_URL_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({})) + + +@dataclass(frozen=True) +class MiniCPM41E2EConfig: + model_path: str + request_model: str + port: int + metrics_port: int + queue_port: int + cache_queue_port: int + max_model_len: str + max_num_seqs: str + tensor_parallel_size: str + startup_timeout: int + log_path: Path + attention_backend: str + quantization: str | None + extra_args: tuple[str, ...] + + @property + def base_url(self) -> str: + return f"http://127.0.0.1:{self.port}" + + @property + def chat_url(self) -> str: + return f"{self.base_url}/v1/chat/completions" + + +def _env_int(name: str, default: int) -> int: + return int(os.getenv(name, str(default))) + + +def _make_config(log_path: Path) -> MiniCPM41E2EConfig: + return MiniCPM41E2EConfig( + model_path=os.getenv("MODEL_PATH") or os.getenv("MINICPM41_MODEL_PATH", "openbmb/MiniCPM4.1-8B"), + request_model=os.getenv("MINICPM41_E2E_REQUEST_MODEL", "default"), + port=_env_int("FD_API_PORT", 8188), + metrics_port=_env_int("FD_METRICS_PORT", 8233), + queue_port=_env_int("FD_ENGINE_QUEUE_PORT", 8133), + cache_queue_port=_env_int("FD_CACHE_QUEUE_PORT", 8333), + max_model_len=os.getenv("MINICPM41_E2E_MAX_MODEL_LEN", "4096"), + max_num_seqs=os.getenv("MINICPM41_E2E_MAX_NUM_SEQS", "2"), + tensor_parallel_size=os.getenv("MINICPM41_E2E_TP", "1"), + startup_timeout=_env_int("MINICPM41_E2E_TIMEOUT", 300), + log_path=log_path, + attention_backend=os.getenv("FD_ATTENTION_BACKEND", "FLASH_ATTN"), + quantization=os.getenv("MINICPM41_E2E_QUANTIZATION") or None, + extra_args=tuple(shlex.split(os.getenv("MINICPM41_E2E_EXTRA_ARGS", ""))), + ) + + +def _use_available_default_ports(config: MiniCPM41E2EConfig) -> MiniCPM41E2EConfig: + """Avoid collisions while preserving explicitly configured E2E ports.""" + port_env_by_field = { + "port": "FD_API_PORT", + "metrics_port": "FD_METRICS_PORT", + "queue_port": "FD_ENGINE_QUEUE_PORT", + "cache_queue_port": "FD_CACHE_QUEUE_PORT", + } + updates: dict[str, int] = {} + reserved_ports = { + getattr(config, field) for field, env_name in port_env_by_field.items() if env_name in os.environ + } + for field, env_name in port_env_by_field.items(): + if env_name in os.environ: + continue + while True: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + candidate = sock.getsockname()[1] + if candidate not in reserved_ports: + reserved_ports.add(candidate) + updates[field] = candidate + break + return replace(config, **updates) + + +def _build_server_command(config: MiniCPM41E2EConfig) -> list[str]: + cmd = [ + sys.executable, + "-m", + "fastdeploy.entrypoints.openai.api_server", + "--model", + config.model_path, + "--port", + str(config.port), + "--metrics-port", + str(config.metrics_port), + "--engine-worker-queue-port", + str(config.queue_port), + "--cache-queue-port", + str(config.cache_queue_port), + "--tensor-parallel-size", + config.tensor_parallel_size, + "--max-model-len", + config.max_model_len, + "--max-num-seqs", + config.max_num_seqs, + ] + if config.quantization: + cmd.extend(["--quantization", config.quantization]) + cmd.extend(config.extra_args) + if config.attention_backend == "INFLLMV2_ATTN": + # InfLLM-V2 rejects both CUDA Graph and prefix caching. Keep the E2E + # command runnable when callers select that backend (as go_run_val.sh + # does) instead of relying on two unrelated defaults. + if "--graph-optimization-config" not in config.extra_args: + cmd.extend(["--graph-optimization-config", '{"use_cudagraph": false}']) + if "--no-enable-prefix-caching" not in config.extra_args: + cmd.append("--no-enable-prefix-caching") + return cmd + + +def test_make_config_reads_quantization_env(tmp_path, monkeypatch): + monkeypatch.setenv("MODEL_PATH", "/models/minicpm41") + monkeypatch.setenv("MINICPM41_E2E_QUANTIZATION", "wint4") + monkeypatch.setenv("MINICPM41_E2E_EXTRA_ARGS", "--served-model-name minicpm41") + monkeypatch.delenv("FD_ATTENTION_BACKEND", raising=False) + + config = _make_config(tmp_path / "server.log") + + assert config.model_path == "/models/minicpm41" + assert config.attention_backend == "FLASH_ATTN" + assert config.quantization == "wint4" + assert config.extra_args == ("--served-model-name", "minicpm41") + + +def test_auto_selected_ports_preserve_explicit_configuration(tmp_path, monkeypatch): + monkeypatch.setenv("FD_API_PORT", "18188") + monkeypatch.delenv("FD_METRICS_PORT", raising=False) + monkeypatch.delenv("FD_ENGINE_QUEUE_PORT", raising=False) + monkeypatch.delenv("FD_CACHE_QUEUE_PORT", raising=False) + + config = _use_available_default_ports(_make_config(tmp_path / "server.log")) + + assert config.port == 18188 + assert len({config.port, config.metrics_port, config.queue_port, config.cache_queue_port}) == 4 + + +@pytest.mark.parametrize("quantization", ["wint4", "wint8"]) +def test_build_server_command_includes_quantization(tmp_path, quantization): + config = MiniCPM41E2EConfig( + model_path="/models/minicpm41", + request_model="default", + port=8188, + metrics_port=8233, + queue_port=8133, + cache_queue_port=8333, + max_model_len="4096", + max_num_seqs="1", + tensor_parallel_size="1", + startup_timeout=300, + log_path=tmp_path / "server.log", + attention_backend="FLASH_ATTN", + quantization=quantization, + extra_args=("--served-model-name", "minicpm41"), + ) + + cmd = _build_server_command(config) + + assert cmd[cmd.index("--quantization") + 1] == quantization + assert cmd[-2:] == ["--served-model-name", "minicpm41"] + + +def test_build_server_command_makes_infllmv2_configuration_compatible(tmp_path): + config = MiniCPM41E2EConfig( + model_path="/models/minicpm41", + request_model="default", + port=8188, + metrics_port=8233, + queue_port=8133, + cache_queue_port=8333, + max_model_len="4096", + max_num_seqs="2", + tensor_parallel_size="1", + startup_timeout=300, + log_path=tmp_path / "server.log", + attention_backend="INFLLMV2_ATTN", + quantization=None, + extra_args=(), + ) + + cmd = _build_server_command(config) + + assert cmd[cmd.index("--graph-optimization-config") + 1] == '{"use_cudagraph": false}' + assert "--no-enable-prefix-caching" in cmd + + +def _read_log(log_path: Path, max_chars: int = 20000) -> str: + if not log_path.exists(): + return "" + content = log_path.read_text(encoding="utf-8", errors="replace") + return content[-max_chars:] + + +def _http_get(url: str, timeout: int = 2) -> tuple[int, str]: + with LOCAL_URL_OPENER.open(url, timeout=timeout) as response: + return response.status, response.read().decode("utf-8") + + +def _http_post_json(url: str, payload: dict, timeout: int = 60) -> tuple[int, dict]: + data = json.dumps(payload).encode("utf-8") + request = urllib.request.Request( + url, + data=data, + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with LOCAL_URL_OPENER.open(request, timeout=timeout) as response: + return response.status, json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + body = exc.read().decode("utf-8", errors="replace") + try: + payload = json.loads(body) + except json.JSONDecodeError: + payload = {"error": body} + return exc.code, payload + + +def _thinking_token_sequences(model_path: str) -> tuple[list[int], list[int]]: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + + def encode(text: str) -> list[int]: + return tokenizer.encode(text, add_special_tokens=False) + + prefix_ids = encode("x") + contextual_start = encode("x") + contextual_end = encode("x\n\n") + assert contextual_start[: len(prefix_ids)] == prefix_ids + assert contextual_end[: len(prefix_ids)] == prefix_ids + return contextual_start[len(prefix_ids) :], contextual_end[len(prefix_ids) :] + + +def _wait_until_healthy(process: subprocess.Popen, config: MiniCPM41E2EConfig) -> None: + deadline = time.time() + config.startup_timeout + health_url = f"{config.base_url}/health" + while time.time() < deadline: + if process.poll() is not None: + raise RuntimeError( + "MiniCPM4.1 server exited before health check passed.\n" + f"Command: {' '.join(_build_server_command(config))}\n" + f"Log tail:\n{_read_log(config.log_path)}" + ) + try: + status, _ = _http_get(health_url) + if status == 200: + return + except (TimeoutError, urllib.error.URLError): + time.sleep(1) + raise RuntimeError( + "MiniCPM4.1 server did not become healthy before timeout.\n" + f"Command: {' '.join(_build_server_command(config))}\n" + f"Log tail:\n{_read_log(config.log_path)}" + ) + + +def _terminate_process(process: subprocess.Popen) -> None: + if process.poll() is not None: + return + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=30) + + +@pytest.fixture(scope="session") +def minicpm41_server(tmp_path_factory): + log_path = tmp_path_factory.mktemp("minicpm41_e2e") / "server.log" + config = _use_available_default_ports(_make_config(log_path)) + env = os.environ.copy() + env["FD_ATTENTION_BACKEND"] = config.attention_backend + env.setdefault("FD_MODEL_SOURCE", "HUGGINGFACE") + # The engine launches worker_process.py as a script. Add the repository to + # PYTHONPATH so source-tree E2E runs do not depend on an editable install. + pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = str(REPO_ROOT) if not pythonpath else f"{REPO_ROOT}{os.pathsep}{pythonpath}" + + cmd = _build_server_command(config) + with config.log_path.open("w", encoding="utf-8") as logfile: + process = subprocess.Popen( + cmd, + env=env, + stdout=logfile, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + + try: + _wait_until_healthy(process, config) + yield config + finally: + _terminate_process(process) + + +def _assert_chat_response(payload: dict) -> None: + assert isinstance(payload.get("id"), str) + assert payload.get("object") == "chat.completion" + assert isinstance(payload.get("choices"), list) + assert payload["choices"] + choice = payload["choices"][0] + assert choice["index"] == 0 + assert choice["message"]["role"] == "assistant" + assert isinstance(choice["message"]["content"], str) + assert choice["message"]["content"] + assert choice.get("finish_reason") in {"stop", "length"} + if "usage" in payload: + assert payload["usage"]["prompt_tokens"] > 0 + assert payload["usage"]["completion_tokens"] > 0 + + +def _assert_non_thinking_response(payload: dict) -> None: + _assert_chat_response(payload) + message = payload["choices"][0]["message"] + assert "" not in message["content"] + assert "" not in message["content"] + assert message.get("reasoning_content") in (None, "") + + +def _assert_forced_end_after_decode_budget(payload: dict, budget: int, forced_end_ids: list[int]) -> None: + """Check the request-prefix/decode-budget contract used by this E2E test. + + Request-side ``completion_token_ids`` are appended to ``prompt_token_ids``. + They can place the state machine inside ````, but prompt-side + reasoning tokens do not consume ``reasoning_max_tokens``. The budget counts + newly decoded reasoning tokens, so the forced ``\n\n`` sequence must + begin at ``budget`` in the response-side completion token ids. + """ + generated_ids = payload["choices"][0]["message"]["completion_token_ids"] + forced_end_slice = generated_ids[budget : budget + len(forced_end_ids)] + assert forced_end_slice == forced_end_ids, generated_ids + + +def _max_tokens_with_forced_end(budget: int, forced_end_ids: list[int]) -> int: + """Leave one slot for the serving layer's terminal EOS token.""" + return budget + len(forced_end_ids) + 1 + + +def test_minicpm41_openai_chat_completion_e2e(minicpm41_server): + status, payload = _http_post_json( + minicpm41_server.chat_url, + { + "model": minicpm41_server.request_model, + "messages": [ + {"role": "system", "content": "You are a concise assistant."}, + {"role": "user", "content": "Say hello in one short sentence."}, + ], + "temperature": 0, + "max_tokens": 16, + }, + ) + + assert status == 200, payload + _assert_chat_response(payload) + + +def test_minicpm41_chat_completion_with_history_e2e(minicpm41_server): + status, payload = _http_post_json( + minicpm41_server.chat_url, + { + "model": minicpm41_server.request_model, + "messages": [ + {"role": "user", "content": "Remember this marker: FD_MINICPM41_E2E."}, + {"role": "assistant", "content": "I will remember FD_MINICPM41_E2E."}, + {"role": "user", "content": "Reply with the marker only."}, + ], + "temperature": 0, + "max_tokens": 16, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + + assert status == 200, payload + _assert_chat_response(payload) + + +def test_minicpm41_forces_multitoken_thinking_end_e2e(minicpm41_server): + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(minicpm41_server.model_path, trust_remote_code=True) + think_start_ids, forced_end_ids = _thinking_token_sequences(minicpm41_server.model_path) + thinking_payload_ids = tokenizer.encode("careful reasoning needs several steps", add_special_tokens=False)[:3] + assert len(thinking_payload_ids) == 3 + reasoning_budget = len(thinking_payload_ids) + + status, payload = _http_post_json( + minicpm41_server.chat_url, + { + "model": minicpm41_server.request_model, + "messages": [{"role": "user", "content": "What is 17 plus 25?"}], + "completion_token_ids": think_start_ids + thinking_payload_ids, + "chat_template_kwargs": {"enable_thinking": True}, + "reasoning_max_tokens": reasoning_budget, + "return_token_ids": True, + "temperature": 0, + "max_tokens": _max_tokens_with_forced_end(reasoning_budget, forced_end_ids), + }, + ) + + assert status == 200, payload + _assert_forced_end_after_decode_budget(payload, reasoning_budget, forced_end_ids) + + +def test_minicpm41_disable_thinking_does_not_emit_think_block_e2e(minicpm41_server): + status, payload = _http_post_json( + minicpm41_server.chat_url, + { + "model": minicpm41_server.request_model, + "messages": [{"role": "user", "content": "Reply with the number 42 only."}], + "chat_template_kwargs": {"enable_thinking": False}, + "temperature": 0, + "max_tokens": 16, + }, + ) + + assert status == 200, payload + _assert_non_thinking_response(payload) + + +def test_minicpm41_mixed_thinking_modes_e2e(minicpm41_server): + """Submit enabled and disabled thinking requests concurrently. + + ``max_num_seqs >= 2`` lets the scheduler place both requests in the same + active batch. The token-level mixed-slot behavior is asserted + deterministically in tests/model_executor/test_minicpm41.py. + """ + if int(minicpm41_server.max_num_seqs) < 2: + pytest.skip("mixed thinking E2E requires MINICPM41_E2E_MAX_NUM_SEQS >= 2") + + think_start_ids, forced_end_ids = _thinking_token_sequences(minicpm41_server.model_path) + thinking_budget = 1 + thinking_payload = { + "model": minicpm41_server.request_model, + "messages": [{"role": "user", "content": "What is 17 plus 25?"}], + "completion_token_ids": think_start_ids, + "chat_template_kwargs": {"enable_thinking": True}, + "reasoning_max_tokens": thinking_budget, + "return_token_ids": True, + "temperature": 0, + "max_tokens": _max_tokens_with_forced_end(thinking_budget, forced_end_ids), + } + non_thinking_payload = { + "model": minicpm41_server.request_model, + "messages": [{"role": "user", "content": "Reply with the number 42 only."}], + "chat_template_kwargs": {"enable_thinking": False}, + "temperature": 0, + "max_tokens": 16, + } + + with ThreadPoolExecutor(max_workers=2) as executor: + thinking_future = executor.submit(_http_post_json, minicpm41_server.chat_url, thinking_payload) + non_thinking_future = executor.submit(_http_post_json, minicpm41_server.chat_url, non_thinking_payload) + thinking_status, thinking_response = thinking_future.result() + non_thinking_status, non_thinking_response = non_thinking_future.result() + + assert thinking_status == 200, thinking_response + assert non_thinking_status == 200, non_thinking_response + _assert_forced_end_after_decode_budget(thinking_response, thinking_budget, forced_end_ids) + _assert_non_thinking_response(non_thinking_response) + + +@pytest.mark.skip(reason="Skipping because the long-context MiniCPM4.1 E2E test requires target GPU memory.") +def test_minicpm41_long_context_chat_completion_e2e(minicpm41_server): + prompt_chars = _env_int("MINICPM41_LONG_PROMPT_CHARS", 20000) + long_context = ("0123456789abcdef " * ((prompt_chars // 17) + 1))[:prompt_chars] + status, payload = _http_post_json( + minicpm41_server.chat_url, + { + "model": minicpm41_server.request_model, + "messages": [ + { + "role": "user", + "content": f"{long_context}\n\nAnswer with one word: ready.", + } + ], + "temperature": 0, + "max_tokens": 16, + }, + timeout=120, + ) + + assert status == 200, payload + _assert_chat_response(payload) diff --git a/tests/engine/test_common_engine.py b/tests/engine/test_common_engine.py index 26fd38bc898..dde0ca6b302 100644 --- a/tests/engine/test_common_engine.py +++ b/tests/engine/test_common_engine.py @@ -37,6 +37,7 @@ from fastdeploy.engine.args_utils import EngineArgs from fastdeploy.engine.common_engine import ( EngineService, + _build_model_thinking_token_sequences, _format_worker_launch_failure_message, _read_latest_worker_traceback, ) @@ -3910,3 +3911,56 @@ def test_format_worker_launch_failure_message_without_traceback(self): self.assertIn("Failed to launch worker processes", result) self.assertIn("workerlog.*", result) self.assertNotIn("Traceback", result) + + +class TestModelThinkingTokenSequenceHook(unittest.TestCase): + @staticmethod + def _model_config(model_cls): + registry = Mock() + registry.resolve_model_cls.return_value = (model_cls, "StubForCausalLM") + return types.SimpleNamespace( + registry=registry, + architectures=["StubForCausalLM"], + vocab_size=100, + _model_info=types.SimpleNamespace(), + ) + + def test_uninspected_model_config_skips_builder_resolution(self): + config = types.SimpleNamespace( + registry=Mock(), + architectures=["StubForCausalLM"], + vocab_size=100, + _model_info=None, + ) + + result = _build_model_thinking_token_sequences(config, Mock()) + + self.assertIsNone(result) + config.registry.resolve_model_cls.assert_not_called() + + def test_model_without_builder_does_not_enable_sequence_mode(self): + class Model: + pass + + result = _build_model_thinking_token_sequences(self._model_config(Model), Mock()) + + self.assertIsNone(result) + + def test_model_builder_result_is_returned(self): + expected = {"start": [[1, 2]], "end": [[3, 4]], "forced_end": [3, 4]} + + class Model: + build_thinking_token_sequences = Mock(return_value=expected) + + tokenizer = Mock() + result = _build_model_thinking_token_sequences(self._model_config(Model), tokenizer) + + self.assertEqual(result, expected) + Model.build_thinking_token_sequences.assert_called_once_with(tokenizer) + + def test_model_builder_failure_propagates(self): + class Model: + build_thinking_token_sequences = Mock(side_effect=ValueError("invalid tokenizer markers")) + + with self.assertRaisesRegex(ValueError, "invalid tokenizer markers"): + _build_model_thinking_token_sequences(self._model_config(Model), Mock()) diff --git a/tests/model_executor/test_infllmv2_attention_backend.py b/tests/model_executor/test_infllmv2_attention_backend.py new file mode 100644 index 00000000000..bbab54233b9 --- /dev/null +++ b/tests/model_executor/test_infllmv2_attention_backend.py @@ -0,0 +1,325 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import weakref +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def read_repo_file(path): + return (REPO_ROOT / path).read_text(encoding="utf-8") + + +def make_config(**sparse_overrides): + sparse = { + "kernel_size": 4, + "kernel_stride": 2, + "topk": 2, + "dense_len": 16, + "init_blocks": 1, + "window_size": 8, + } + sparse.update(sparse_overrides) + return SimpleNamespace( + cache_config=SimpleNamespace( + block_size=8, + enable_prefix_caching=False, + num_cpu_blocks=0, + kvcache_storage_backend=None, + ), + model_config=SimpleNamespace( + sparse_config=sparse, + max_model_len=128, + causal=True, + head_dim=8, + num_hidden_layers=2, + num_key_value_heads=2, + start_layer_index=0, + ), + speculative_config=SimpleNamespace(method=None, num_speculative_tokens=0, model_type="main"), + parallel_config=SimpleNamespace( + pd_disaggregation_mode=None, + local_data_parallel_id=0, + tensor_parallel_size=1, + tensor_parallel_rank=0, + ), + scheduler_config=SimpleNamespace(splitwise_role="mixed"), + graph_opt_config=SimpleNamespace(use_cudagraph=False, full_cuda_graph=True), + enable_rope_3d_runtime=False, + ) + + +def make_backend(**sparse_overrides): + paddle = pytest.importorskip("paddle") + from fastdeploy.platforms import current_platform + + with patch.object(current_platform, "is_cuda", return_value=False): + from fastdeploy.model_executor.layers.attention.infllmv2_attention_backend import ( + InfLLMV2AttentionBackend, + ) + + backend = InfLLMV2AttentionBackend( + fd_config=make_config(**sparse_overrides), + kv_num_heads=2, + num_heads=4, + head_dim=8, + ) + forward_meta = SimpleNamespace() + backend.init_attention_metadata(forward_meta) + return backend, forward_meta, paddle + + +def test_infllmv2_backend_registration_and_export(): + from fastdeploy.platforms.base import _Backend + from fastdeploy.platforms.cuda import CUDAPlatform + + assert _Backend.INFLLMV2_ATTN.name == "INFLLMV2_ATTN" + assert ( + CUDAPlatform.get_attention_backend_cls(_Backend.INFLLMV2_ATTN) + == "fastdeploy.model_executor.layers.attention.InfLLMV2AttentionBackend" + ) + + +def test_infllmv2_metadata_uses_sparse_config_and_fixed_capacity(): + backend, forward_meta, _ = make_backend() + + metadata = backend.attention_metadata + assert metadata.kernel_size == 4 + assert metadata.kernel_stride == 2 + assert metadata.topk == 2 + assert metadata.dense_len == 16 + assert metadata.block_size == 8 + assert metadata.init_blocks == 1 + assert metadata.local_blocks == 1 + assert metadata.selected_capacity == 3 + assert forward_meta.attn_metadata is metadata + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"kernel_stride": 3}, "divisible by kernel_stride"), + ({"topk": 0}, "topk must be positive"), + ({"dense_len": 15}, "coarse semantic window"), + ({"init_blocks": 2}, "smaller than topk"), + ({"window_size": 7}, "multiple of block_size"), + ({"prefill_query_chunk_size": 64}, "positive multiple of 128"), + ], +) +def test_infllmv2_rejects_invalid_sparse_config(overrides, message): + paddle = pytest.importorskip("paddle") + from fastdeploy.platforms import current_platform + + with patch.object(current_platform, "is_cuda", return_value=False): + from fastdeploy.model_executor.layers.attention.infllmv2_attention_backend import ( + InfLLMV2AttentionBackend, + ) + + with pytest.raises(ValueError, match=message): + InfLLMV2AttentionBackend(make_config(**overrides), 2, 4, 8) + assert paddle is not None + + +def test_infllmv2_split_qkv_matches_fastdeploy_layout(): + backend, _, paddle = make_backend() + qkv = paddle.arange(4 * (4 + 2 + 2) * 8, dtype="float32").reshape([4, 64]) + + q, k, v = backend._split_qkv(qkv) + + assert list(q.shape) == [4, 4, 8] + assert list(k.shape) == [4, 2, 8] + assert list(v.shape) == [4, 2, 8] + assert paddle.allclose(q.reshape([4, 32]), qkv[:, :32]) + assert paddle.allclose(k.reshape([4, 16]), qkv[:, 32:48]) + assert paddle.allclose(v.reshape([4, 16]), qkv[:, 48:]) + + +def test_infllmv2_rejects_bad_qkv_width_and_quantized_cache(): + backend, _, paddle = make_backend() + + with pytest.raises(ValueError, match="fused qkv last dimension"): + backend._split_qkv(paddle.zeros([2, 63], dtype="float32")) + with pytest.raises(ValueError, match="unquantized"): + backend.get_kv_cache_shape(4, "int4_zp") + + +def test_infllmv2_semantic_cache_and_workspace_are_persistent(): + backend, _, paddle = make_backend() + cache = paddle.zeros([6, 2, 8, 8], dtype="float32") + + fine, coarse = backend._ensure_compressed_cache(cache) + assert list(fine.shape) == [6, 2, 4, 8] + assert list(coarse.shape) == [6, 2, 1, 8] + cached_fine, cached_coarse = backend._ensure_compressed_cache(cache) + assert cached_fine is fine + assert cached_coarse is coarse + + backend._ensure_workspace(2, 16, paddle.float32) + first_topk = backend._topk_indices_ws + first_output = backend._attention_out_ws + backend._ensure_workspace(2, 16, paddle.float32) + assert backend._topk_indices_ws is first_topk + assert backend._attention_out_ws is first_output + assert list(first_topk.shape) == [2, 2, 3] + assert list(first_output.shape) == [2, 4, 8] + + backend._release_workspace() + backend._ensure_workspace(2, 16, paddle.float32, allocate_attention=False) + assert backend._attention_out_ws is None + assert backend._partial_acc_ws is None + + +def test_infllmv2_reset_runtime_cache_releases_semantic_summaries(): + backend, _, paddle = make_backend() + backend._ensure_compressed_cache(paddle.zeros([4, 2, 8, 8])) + + backend.reset_runtime_cache() + + assert backend._compressed_k is None + assert backend._compressed_k2 is None + assert backend.attention_metadata.compressed_k is None + assert backend.attention_metadata.compressed_k2 is None + + +def test_infllmv2_sparse_decode_requires_prefill_semantic_cache(): + backend, forward_meta, paddle = make_backend() + forward_meta.caches = [paddle.zeros([2, 2, 8, 8]), paddle.zeros([2, 2, 8, 8])] + + with pytest.raises(RuntimeError, match="initialized by prefill"): + backend._forward_sparse_decode(paddle.zeros([1, 64]), SimpleNamespace(layer_id=0), forward_meta) + + +def test_infllmv2_decode_orders_writer_update_stage1_and_stage2(): + backend, forward_meta, paddle = make_backend() + cache_k = paddle.zeros([4, 2, 8, 8]) + cache_v = paddle.zeros_like(cache_k) + backend._compressed_k = paddle.zeros([4, 2, 4, 8]) + backend._compressed_k2 = paddle.zeros([4, 2, 1, 8]) + backend._compressed_cache_owner = weakref.ref(cache_k) + forward_meta.caches = [cache_k, cache_v] + forward_meta.block_tables = paddle.arange(4, dtype="int32").reshape([1, 4]) + forward_meta.seq_lens_decoder = paddle.to_tensor([31], dtype="int32") + forward_meta.seq_lens_this_time = paddle.ones([1], dtype="int32") + forward_meta.batch_id_per_token = paddle.zeros([1], dtype="int32") + forward_meta.cu_seqlens_q = paddle.to_tensor([0, 1], dtype="int32") + raw_qkv = paddle.zeros([1, 64]) + post_rope_qkv = paddle.arange(64, dtype="float32").reshape([1, 64]) + calls = [] + + def select_blocks(query, *args): + calls.append(("stage1", query.clone())) + return ( + backend._topk_indices_ws, + backend._block_scores_ws, + backend._selected_counts_ws, + backend._coarse_lse_ws, + backend._coarse_partial_max_ws, + backend._coarse_partial_sum_ws, + ) + + def sparse_attention(query, *args): + calls.append(("stage2", query.clone())) + return ( + backend._attention_out_ws, + backend._partial_acc_ws, + backend._partial_max_ws, + backend._partial_sum_ws, + ) + + with ( + patch.object(backend, "_prepare_sparse_runtime", side_effect=lambda *args: calls.append(("prepare", None))), + patch.object( + backend, "_write_decode_cache", side_effect=lambda *args: calls.append(("writer", None)) or post_rope_qkv + ), + patch.object(backend, "_update_compressed_cache", side_effect=lambda *args: calls.append(("update", None))), + patch.object(backend, "_load_sparse_ops", return_value=(object(), select_blocks, sparse_attention)), + ): + output = backend._forward_sparse_decode(raw_qkv, SimpleNamespace(layer_id=0), forward_meta) + + assert [name for name, _ in calls] == ["prepare", "writer", "update", "stage1", "stage2"] + expected_query = post_rope_qkv[:, :32].reshape([1, 4, 8]) + assert paddle.allclose(calls[3][1], expected_query) + assert paddle.allclose(calls[4][1], expected_query) + assert list(output.shape) == [1, 32] + + +def test_infllmv2_forward_extend_routes_eligible_prefill_to_sparse_path(): + backend, forward_meta, paddle = make_backend() + qkv = paddle.zeros([24, 64]) + expected = paddle.ones([24, 32]) + layer = SimpleNamespace() + + with ( + patch.object(backend, "_can_use_sparse_prefill", return_value=True), + patch.object(backend, "_forward_sparse_prefill", return_value=expected) as sparse_prefill, + patch.object(backend, "_forward_dense_and_update") as dense_prefill, + ): + output = backend.forward_extend(None, None, None, qkv, None, None, layer, forward_meta) + + assert output is expected + sparse_prefill.assert_called_once_with(qkv, layer, forward_meta) + dense_prefill.assert_not_called() + + +def test_infllmv2_forward_mixed_routes_eligible_initial_prefill_to_sparse_path(): + backend, forward_meta, paddle = make_backend() + forward_meta.exist_prefill = True + qkv = paddle.zeros([24, 64]) + expected = paddle.ones([24, 32]) + layer = SimpleNamespace() + + with ( + patch.object(backend, "_can_use_sparse_prefill", return_value=True), + patch.object(backend, "_forward_sparse_prefill", return_value=expected) as sparse_prefill, + patch.object(backend, "_forward_dense_and_update") as dense_prefill, + ): + output = backend.forward_mixed(None, None, None, qkv, None, None, layer, forward_meta) + + assert output is expected + sparse_prefill.assert_called_once_with(qkv, layer, forward_meta) + dense_prefill.assert_not_called() + + +def test_infllmv2_sparse_prefill_requires_a_full_history_selection(): + backend, _, paddle = make_backend() + backend.sparse_prefill = True + backend.causal = True + backend.block_size = 64 + backend.head_dim = 128 + backend.num_heads = 32 + backend.kv_num_heads = 2 + backend.dense_len = 8192 + backend.topk = 1 + backend.local_blocks = 1 + qkv = paddle.zeros([8193, 1], dtype="bfloat16") + forward_meta = SimpleNamespace( + block_tables=paddle.zeros([1, 129], dtype="int32"), + attn_mask_offsets=None, + max_len_tensor_cpu=paddle.zeros([3], dtype="int32"), + ) + + assert not backend._can_use_sparse_prefill(qkv, forward_meta) + backend.topk = 64 + backend.local_blocks = 32 + assert backend._can_use_sparse_prefill(qkv, forward_meta) + + +def test_infllmv2_additional_cache_cost_matches_two_summary_scales(): + backend, _, _ = make_backend() + assert backend.get_additional_cache_block_bytes(2) == 2 * 2 * 8 * (4 + 1) diff --git a/tests/model_executor/test_minicpm41.py b/tests/model_executor/test_minicpm41.py new file mode 100644 index 00000000000..474ba4393e6 --- /dev/null +++ b/tests/model_executor/test_minicpm41.py @@ -0,0 +1,520 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import importlib.util +import math +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +MINICPM41_DIR = REPO_ROOT / "fastdeploy/model_executor/models/minicpm41" + +THINK_START_SEQUENCES = [[10, 11], [12, 11]] +THINK_END_SEQUENCES = [[20, 11], [21, 11]] +THINK_FORCED_END_IDS = [30, 21, 11, 31] + + +class Config: + pass + + +def make_hybrid_config(max_thinking_length=2): + model_config = SimpleNamespace( + dtype="float32", + vocab_size=40, + think_start_id=-1, + think_end_id=-1, + line_break_id=-1, + think_token_sequences={ + "start": THINK_START_SEQUENCES, + "end": THINK_END_SEQUENCES, + "forced_end": THINK_FORCED_END_IDS, + }, + reasoning_tokens="thinking", + max_thinking_length=max_thinking_length, + ) + return SimpleNamespace(model_config=model_config) + + +def make_hybrid_share_inputs( + prompt_ids, *, budget=-1, enable_thinking=True, req_id="req-1", logits_processors_args=None +): + import paddle + + padded_prompt = list(prompt_ids) + [-1] * 16 + return { + "stop_flags": paddle.to_tensor([[False]], dtype="bool"), + "enable_thinking": paddle.to_tensor([[enable_thinking]], dtype="bool"), + "max_think_lens": paddle.to_tensor([[budget]], dtype="int32"), + "prompt_lens": paddle.to_tensor([[len(prompt_ids)]], dtype="int64"), + "step_idx": paddle.to_tensor([[0]], dtype="int64"), + "next_tokens": paddle.to_tensor([[-1]], dtype="int64"), + "req_ids": [req_id], + "logits_processors_args": [logits_processors_args or {}], + "prompt_ids": paddle.to_tensor([padded_prompt], dtype="int64"), + "token_ids_all": paddle.to_tensor([padded_prompt], dtype="int64"), + "pre_ids": paddle.to_tensor([[-1] * 16], dtype="int64"), + } + + +def run_hybrid_step(mode, share_inputs, step_idx, next_token, vocab_size=40): + import paddle + + share_inputs["step_idx"][0, 0] = step_idx + share_inputs["next_tokens"][0, 0] = next_token + mode.update_state(share_inputs) + logits = paddle.ones([len(share_inputs["req_ids"]), vocab_size], dtype="float32") + return mode.apply(logits) + + +def finite_token_ids(logits, slot_id=0): + import paddle + + return paddle.nonzero(~paddle.isinf(logits[slot_id])).reshape([-1]).numpy().tolist() + + +def read_repo_file(path): + return (REPO_ROOT / path).read_text(encoding="utf-8") + + +def read_model_source(): + return read_repo_file("fastdeploy/model_executor/models/minicpm41/minicpm41.py") + + +def load_module(module_path): + spec = importlib.util.spec_from_file_location(module_path.stem, module_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def load_hybrid_reasoning_module(): + package_name = "_test_minicpm41_pkg" + package = types.ModuleType(package_name) + package.__path__ = [str(MINICPM41_DIR)] + sys.modules[package_name] = package + + spec = importlib.util.spec_from_file_location( + f"{package_name}.hybrid_reasoning", + MINICPM41_DIR / "hybrid_reasoning.py", + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def load_minicpm41_model_module(): + from fastdeploy.platforms import current_platform + + with patch.object(current_platform, "is_cuda", return_value=False): + return importlib.import_module("fastdeploy.model_executor.models.minicpm41.minicpm41") + + +def test_minicpm41_model_package_exists_for_auto_registry(): + assert (REPO_ROOT / "fastdeploy/model_executor/models/minicpm41/__init__.py").exists() + assert (REPO_ROOT / "fastdeploy/model_executor/models/minicpm41/minicpm41.py").exists() + + +def test_minicpm41_causallm_registers_hf_architecture_name(): + module = load_minicpm41_model_module() + from fastdeploy.model_executor.models.model_base import ModelCategory, ModelRegistry + + assert module.MiniCPM41ForCausalLM.name() == "MiniCPMForCausalLM" + assert module.MiniCPM41PretrainedModel.arch_name() == "MiniCPMForCausalLM" + assert ModelRegistry._arch_to_model_cls["MiniCPMForCausalLM"] is module.MiniCPM41ForCausalLM + assert ModelRegistry._enhanced_models["MiniCPMForCausalLM"]["module_name"] == "minicpm41.minicpm41" + category = ModelRegistry._enhanced_models["MiniCPMForCausalLM"]["category"] + assert ModelCategory.TEXT_GENERATION in category + assert ModelCategory.REASONING in category + assert ModelRegistry().is_reasoning_model("MiniCPMForCausalLM") + model_cls, _ = ModelRegistry().resolve_model_cls("MiniCPMForCausalLM") + assert hasattr(model_cls, "build_thinking_token_sequences") + + +def test_minicpm41_model_uses_functional_decoder_components(): + source = read_model_source() + + assert "class MiniCPM41Model" in source + assert "class MiniCPM41DecoderLayer" in source + assert "class MiniCPM41Attention" in source + assert "class MiniCPM41MLP" in source + assert "self.minicpm41 = MiniCPM41Model" in source + assert "NotImplementedError" not in source + + +def test_minicpm41_qkv_projection_is_bias_free_for_hf_weights(): + source = read_model_source() + + assert "QKVParallelLinear" in source + assert "with_bias=False" in source + + +def test_minicpm41_longrope_matches_reference_frequency_and_neox_layout(): + import paddle + + from fastdeploy.model_executor.layers.rotary_embedding import get_rope_impl + + config = Config() + config.architectures = ["MiniCPMForCausalLM"] + config.max_position_embeddings = 8 + config.rope_scaling = { + "rope_type": "longrope", + "short_factor": [1.0, 2.0], + "long_factor": [2.0, 4.0], + "original_max_position_embeddings": 4, + } + position_ids = paddle.arange(4).reshape([1, -1]) + + actual = get_rope_impl(4, 100.0, position_ids, config) + + inv_freq = paddle.to_tensor([1.0, 0.1], dtype="float32") / paddle.to_tensor([1.0, 2.0]) + freqs = position_ids.cast("float32").unsqueeze(-1) * inv_freq.reshape([1, 1, -1]) + emb = paddle.concat([freqs, freqs], axis=-1).reshape([1, 4, 1, 4]) + magnitude_scale = math.sqrt(1 + math.log(2) / math.log(4)) + expected = paddle.stack([paddle.cos(emb), paddle.sin(emb)], axis=0) * magnitude_scale + + assert list(actual.shape) == [2, 1, 4, 1, 4] + assert paddle.allclose(actual, expected) + + +def test_minicpm41_weight_mapping_matches_hf_prefixes_and_stacked_weights(): + source = read_model_source() + + assert 'WeightsMapper(orig_to_new_prefix={"model.": "minicpm41."})' in source + assert '("qkv_proj", "q_proj", "q")' in source + assert '("qkv_proj", "k_proj", "k")' in source + assert '("qkv_proj", "v_proj", "v")' in source + assert '("up_gate_proj", "gate_proj", "gate")' in source + assert '("up_gate_proj", "up_proj", "up")' in source + + +def test_minicpm41_qkv_uses_standard_online_quantization_loader(): + source = read_model_source() + + assert "self.qkv_proj = QKVParallelLinear(" in source + assert "MiniCPM41QKVParallelLinear" not in source + assert "load_minicpm41_wint4_qkv_weight" not in source + + +def test_minicpm41_model_uses_mup_scaling_points(): + module = load_minicpm41_model_module() + config = Config() + config.scale_emb = 12 + config.scale_depth = 1.4 + config.num_hidden_layers = 32 + config.hidden_size = 4096 + config.dim_model_base = 256 + + assert module.minicpm41_embedding_scale(config) == 12.0 + assert round(module.minicpm41_residual_scale(config), 8) == 0.24748737 + assert module.minicpm41_lm_head_scale(config) == 0.0625 + + +def test_minicpm41_scaling_helpers_read_pretrained_config_fallbacks(): + module = load_minicpm41_model_module() + pretrained_config = Config() + pretrained_config.scale_emb = 8 + pretrained_config.scale_depth = 2.0 + pretrained_config.num_hidden_layers = 16 + pretrained_config.hidden_size = 1024 + pretrained_config.dim_model_base = 256 + model_config = Config() + model_config.pretrained_config = pretrained_config + + assert module.minicpm41_embedding_scale(model_config) == 8.0 + assert module.minicpm41_residual_scale(model_config) == 0.5 + assert module.minicpm41_lm_head_scale(model_config) == 0.25 + + +def test_hybrid_reasoning_reads_top_level_config_overrides(): + module = load_hybrid_reasoning_module() + fd_config = make_hybrid_config() + fd_config.reasoning_tokens = "analysis" + fd_config.max_thinking_length = 64 + + mode = module.HybridReasoningMode(fd_config) + + assert mode.reasoning_tokens == "analysis" + assert mode.max_thinking_length == 64 + assert mode.think_start_sequences == ((10, 11), (12, 11)) + assert mode.think_end_sequences == ((20, 11), (21, 11)) + assert mode.think_forced_end_ids == [30, 21, 11, 31] + assert mode._sequence_mode + assert mode._enabled + + +def test_hybrid_reasoning_uses_rfc_defaults(): + module = load_hybrid_reasoning_module() + fd_config = make_hybrid_config() + del fd_config.model_config.reasoning_tokens + del fd_config.model_config.max_thinking_length + + mode = module.HybridReasoningMode(fd_config) + + assert mode.reasoning_tokens == "thinking" + assert mode.max_thinking_length == 512 + + +def test_hybrid_reasoning_forces_complete_multitoken_end_at_budget(): + module = load_hybrid_reasoning_module() + mode = module.HybridReasoningMode(make_hybrid_config(max_thinking_length=2)) + share_inputs = make_hybrid_share_inputs([10, 11]) + + assert len(finite_token_ids(run_hybrid_step(mode, share_inputs, 0, -1))) == 40 + assert len(finite_token_ids(run_hybrid_step(mode, share_inputs, 1, 5))) == 40 + assert finite_token_ids(run_hybrid_step(mode, share_inputs, 2, 6)) == [30] + assert finite_token_ids(run_hybrid_step(mode, share_inputs, 3, 30)) == [21] + assert finite_token_ids(run_hybrid_step(mode, share_inputs, 4, 21)) == [11] + assert finite_token_ids(run_hybrid_step(mode, share_inputs, 5, 11)) == [31] + assert len(finite_token_ids(run_hybrid_step(mode, share_inputs, 6, 31))) == 40 + + +def test_hybrid_reasoning_natural_end_and_new_round_reset_state(): + module = load_hybrid_reasoning_module() + mode = module.HybridReasoningMode(make_hybrid_config(max_thinking_length=8)) + share_inputs = make_hybrid_share_inputs([10, 11]) + + run_hybrid_step(mode, share_inputs, 0, -1) + run_hybrid_step(mode, share_inputs, 1, 5) + run_hybrid_step(mode, share_inputs, 2, 20) + natural_end_logits = run_hybrid_step(mode, share_inputs, 3, 11) + + assert len(finite_token_ids(natural_end_logits)) == 40 + assert mode._states["req-1"].ended + + run_hybrid_step(mode, share_inputs, 4, 12) + run_hybrid_step(mode, share_inputs, 5, 11) + run_hybrid_step(mode, share_inputs, 6, 7) + + assert mode._states["req-1"].started + assert not mode._states["req-1"].ended + assert mode._states["req-1"].tokens_after_start == 1 + + +def test_hybrid_reasoning_disabled_request_does_not_change_logits(): + module = load_hybrid_reasoning_module() + mode = module.HybridReasoningMode(make_hybrid_config(max_thinking_length=1)) + share_inputs = make_hybrid_share_inputs( + [10, 11], enable_thinking=False, logits_processors_args={"thinking_budget": 0} + ) + + logits = run_hybrid_step(mode, share_inputs, 1, 5) + + assert len(finite_token_ids(logits)) == 40 + assert mode._states == {} + + +def test_hybrid_reasoning_mixed_batch_only_limits_enabled_request(): + import paddle + + module = load_hybrid_reasoning_module() + mode = module.HybridReasoningMode(make_hybrid_config(max_thinking_length=1)) + share_inputs = { + "stop_flags": paddle.to_tensor([[False], [False]], dtype="bool"), + "enable_thinking": paddle.to_tensor([[True], [False]], dtype="bool"), + "max_think_lens": paddle.to_tensor([[-1], [-1]], dtype="int32"), + "prompt_lens": paddle.to_tensor([[2], [2]], dtype="int64"), + "step_idx": paddle.to_tensor([[0], [0]], dtype="int64"), + "next_tokens": paddle.to_tensor([[-1], [-1]], dtype="int64"), + "req_ids": ["req-thinking", "req-non-thinking"], + # Give the disabled request an explicit zero budget to prove that the + # per-request enable_thinking gate wins before budget validation. + "logits_processors_args": [{}, {"thinking_budget": 0}], + "prompt_ids": paddle.to_tensor([[10, 11] + [-1] * 8, [10, 11] + [-1] * 8], dtype="int64"), + "token_ids_all": paddle.to_tensor([[10, 11] + [-1] * 8, [10, 11] + [-1] * 8], dtype="int64"), + "pre_ids": paddle.to_tensor([[-1] * 10, [-1] * 10], dtype="int64"), + } + + mode.update_state(share_inputs) + share_inputs["step_idx"][:] = 1 + share_inputs["next_tokens"][:] = paddle.to_tensor([[5], [6]], dtype="int64") + mode.update_state(share_inputs) + logits = mode.apply(paddle.ones([2, 40], dtype="float32")) + + assert finite_token_ids(logits, 0) == [30] + assert len(finite_token_ids(logits, 1)) == 40 + assert set(mode._states) == {"req-thinking"} + + +def test_hybrid_reasoning_explicit_budget_overrides_reasoning_max_tokens(): + module = load_hybrid_reasoning_module() + mode = module.HybridReasoningMode(make_hybrid_config(max_thinking_length=8)) + share_inputs = make_hybrid_share_inputs([10, 11], budget=6, logits_processors_args={"thinking_budget": 1}) + + run_hybrid_step(mode, share_inputs, 0, -1) + logits = run_hybrid_step(mode, share_inputs, 1, 5) + + assert finite_token_ids(logits) == [30] + + +def test_hybrid_reasoning_disable_cleans_existing_request_state(): + module = load_hybrid_reasoning_module() + mode = module.HybridReasoningMode(make_hybrid_config(max_thinking_length=4)) + share_inputs = make_hybrid_share_inputs([10, 11]) + run_hybrid_step(mode, share_inputs, 0, -1) + assert "req-1" in mode._states + + share_inputs["enable_thinking"][0, 0] = False + mode.update_state(share_inputs) + + assert "req-1" not in mode._states + + +def test_hybrid_reasoning_rejects_invalid_config_and_budget(): + module = load_hybrid_reasoning_module() + + invalid_max_length = make_hybrid_config() + invalid_max_length.model_config.max_thinking_length = 1.5 + with pytest.raises(ValueError, match="max_thinking_length"): + module.HybridReasoningMode(invalid_max_length) + + malformed_sequences = make_hybrid_config() + malformed_sequences.model_config.think_token_sequences = {"start": [[10, -11]], "end": [[20, 11]]} + with pytest.raises(ValueError, match="think_token_sequences"): + module.HybridReasoningMode(malformed_sequences) + + mode = module.HybridReasoningMode(make_hybrid_config()) + share_inputs = make_hybrid_share_inputs([10, 11], budget=0) + + with pytest.raises(ValueError, match="reasoning_max_tokens"): + mode.update_state(share_inputs) + + +def test_hybrid_reasoning_rejects_missing_marker_sequences(): + module = load_hybrid_reasoning_module() + fd_config = make_hybrid_config() + fd_config.model_config.think_token_sequences = None + + with pytest.raises(ValueError, match="requires valid"): + module.HybridReasoningMode(fd_config) + + +def test_hybrid_reasoning_tracks_request_identity_across_slot_reorder(): + import paddle + + module = load_hybrid_reasoning_module() + mode = module.HybridReasoningMode(make_hybrid_config(max_thinking_length=4)) + share_inputs = { + "stop_flags": paddle.to_tensor([[False], [False]], dtype="bool"), + "enable_thinking": paddle.to_tensor([[True], [True]], dtype="bool"), + "max_think_lens": paddle.to_tensor([[1], [3]], dtype="int32"), + "prompt_lens": paddle.to_tensor([[2], [2]], dtype="int64"), + "step_idx": paddle.to_tensor([[0], [0]], dtype="int64"), + "next_tokens": paddle.to_tensor([[-1], [-1]], dtype="int64"), + "req_ids": ["req-1", "req-2"], + "logits_processors_args": [{}, {}], + "prompt_ids": paddle.to_tensor([[10, 11] + [-1] * 8, [12, 11] + [-1] * 8], dtype="int64"), + "token_ids_all": paddle.to_tensor([[10, 11] + [-1] * 8, [12, 11] + [-1] * 8], dtype="int64"), + "pre_ids": paddle.to_tensor([[-1] * 10, [-1] * 10], dtype="int64"), + } + mode.update_state(share_inputs) + share_inputs["step_idx"][:] = 1 + share_inputs["next_tokens"][:] = paddle.to_tensor([[5], [6]], dtype="int64") + mode.update_state(share_inputs) + first_logits = mode.apply(paddle.ones([2, 40], dtype="float32")) + + assert finite_token_ids(first_logits, 0) == [30] + assert len(finite_token_ids(first_logits, 1)) == 40 + + for key in ( + "stop_flags", + "enable_thinking", + "max_think_lens", + "prompt_lens", + "step_idx", + "next_tokens", + "prompt_ids", + "token_ids_all", + "pre_ids", + ): + share_inputs[key] = paddle.flip(share_inputs[key], axis=[0]) + share_inputs["req_ids"] = ["req-2", "req-1"] + share_inputs["logits_processors_args"] = [{}, {}] + share_inputs["step_idx"][:] = 2 + share_inputs["next_tokens"][:] = paddle.to_tensor([[7], [30]], dtype="int64") + mode.update_state(share_inputs) + reordered_logits = mode.apply(paddle.ones([2, 40], dtype="float32")) + + assert len(finite_token_ids(reordered_logits, 0)) == 40 + assert finite_token_ids(reordered_logits, 1) == [21] + + share_inputs["stop_flags"][1, 0] = True + mode.update_state(share_inputs) + assert "req-1" not in mode._states + + +def test_minicpm41_exposes_owned_hybrid_reasoning_processor(): + model_module = load_minicpm41_model_module() + processor = object() + model = SimpleNamespace(hybrid_reasoning=processor) + + processors = model_module.MiniCPM41ForCausalLM.get_logits_processors(model) + + assert processors == [processor] + + +def test_thinking_token_sequence_builder_keeps_contextual_variants(): + module = load_hybrid_reasoning_module() + encoded = { + "": [1, 2], + "x": [9], + "x": [9, 3, 2], + "": [4, 2], + "x": [9, 5, 2], + "x\n\n": [9, 6, 5, 2, 6], + } + tokenizer = SimpleNamespace(encode=lambda text, add_special_tokens=False: encoded[text]) + + sequences = module.build_minicpm41_thinking_token_sequences(tokenizer) + + assert sequences == { + "start": [[1, 2], [3, 2]], + "end": [[4, 2], [5, 2]], + "forced_end": [6, 5, 2, 6], + } + + +def test_model_runner_registers_model_owned_logits_processor_once(): + from fastdeploy.worker.model_runner_base import ModelRunnerBase + + processor = object() + model = SimpleNamespace(get_logits_processors=lambda: [processor]) + runner = SimpleNamespace( + get_model=lambda: model, + share_inputs={"logits_processors": []}, + ) + + ModelRunnerBase.register_model_logits_processors(runner) + ModelRunnerBase.register_model_logits_processors(runner) + + assert runner.share_inputs["logits_processors"] == [processor] + + +def test_model_runner_rejects_invalid_logits_processor_provider_contract(): + from fastdeploy.worker.model_runner_base import ModelRunnerBase + + model = SimpleNamespace(get_logits_processors=lambda: object()) + runner = SimpleNamespace( + get_model=lambda: model, + share_inputs={"logits_processors": []}, + ) + + with pytest.raises(TypeError, match="must return a list or tuple"): + ModelRunnerBase.register_model_logits_processors(runner) diff --git a/tests/model_executor/test_thinking_budget.py b/tests/model_executor/test_thinking_budget.py index f93afd4d578..90b8bbebf13 100644 --- a/tests/model_executor/test_thinking_budget.py +++ b/tests/model_executor/test_thinking_budget.py @@ -129,6 +129,7 @@ def setUp(self): mock_model_config.think_start_id = THINKING_START_TOKEN_ID mock_model_config.think_end_id = THINKING_END_TOKEN_ID mock_model_config.line_break_id = NEW_LINE_TOKEN_ID + mock_model_config.think_token_sequences = None cache_config = CacheConfig(args_dict_from_engine_args) parallel_config = ParallelConfig(args_dict_from_engine_args) @@ -672,6 +673,300 @@ def test_multiple_requests(self): self.assertEqual(paddle.argmax(processed_batch_logits[2], axis=-1).item(), THINKING_END_TOKEN_ID) +class TestThinkingBudgetTokenSequences(unittest.TestCase): + """Multi-token marker mode for tokenizers where / are sequences.""" + + START_SEQUENCES = [[10, 11], [12, 11]] + END_SEQUENCES = [[20, 11], [21, 11]] + FORCED_END_IDS = [30, 21, 11, 31] + SEQ_VOCAB_SIZE = 40 + + def setUp(self): + engine_args = EngineArgs(max_num_seqs=4) + args_dict_from_engine_args = asdict(engine_args) + + self._fdconfig_patches = [ + patch.object(FDConfig, "read_from_config", return_value=None), + patch.object(FDConfig, "postprocess", return_value=None), + patch.object(FDConfig, "init_pd_info", return_value=None), + patch.object(FDConfig, "check", return_value=None), + ] + for patcher in self._fdconfig_patches: + patcher.start() + self.addCleanup(patcher.stop) + + mock_model_config = MagicMock() + mock_model_config.dtype = "float32" + mock_model_config.vocab_size = self.SEQ_VOCAB_SIZE + mock_model_config.paddle_dtype = paddle.float32 + mock_model_config.max_model_len = 512 + mock_model_config.think_start_id = -1 + mock_model_config.think_end_id = -1 + mock_model_config.line_break_id = -1 + mock_model_config.think_token_sequences = { + "start": self.START_SEQUENCES, + "end": self.END_SEQUENCES, + "forced_end": self.FORCED_END_IDS, + } + + self.fd_config = FDConfig( + model_config=mock_model_config, + cache_config=CacheConfig(args_dict_from_engine_args), + parallel_config=ParallelConfig(args_dict_from_engine_args), + speculative_config=SpeculativeConfig(args_dict_from_engine_args), + scheduler_config=SchedulerConfig(args_dict_from_engine_args), + load_config=LoadConfig(args_dict_from_engine_args), + graph_opt_config=GraphOptimizationConfig(args_dict_from_engine_args), + structured_outputs_config=StructuredOutputsConfig(args_dict_from_engine_args), + router_config=MagicMock(), + test_mode=True, + ) + + def _get_initial_logits(self, batch_size): + logits = paddle.full((batch_size, self.SEQ_VOCAB_SIZE), -10.0, dtype=paddle.float32) + logits[:, 0] = 0.0 + return logits + + def _make_processor(self): + return ThinkingBudgetLogitsProcessor(self.fd_config) + + def test_sequence_mode_enabled(self): + processor = self._make_processor() + self.assertTrue(processor._enabled) + self.assertTrue(processor._sequence_mode) + self.assertEqual(processor.think_start_sequences, ((10, 11), (12, 11))) + self.assertEqual(processor.think_end_sequences, ((20, 11), (21, 11))) + self.assertEqual(processor.think_forced_end_ids, [30, 21, 11, 31]) + + def test_sequence_mode_rejects_non_dict_config(self): + self.fd_config.model_config.think_token_sequences = [] + + with self.assertRaisesRegex(ValueError, "must be a dict"): + self._make_processor() + + def test_sequence_mode_requires_exact_keys(self): + self.fd_config.model_config.think_token_sequences = { + "start": self.START_SEQUENCES, + "end": self.END_SEQUENCES, + } + + with self.assertRaisesRegex(ValueError, "exactly start, end, and forced_end"): + self._make_processor() + + def test_sequence_mode_rejects_invalid_and_out_of_vocab_ids(self): + self.fd_config.model_config.think_token_sequences = { + "start": [[10, -11]], + "end": self.END_SEQUENCES, + "forced_end": self.FORCED_END_IDS, + } + with self.assertRaisesRegex(ValueError, "vocabulary range"): + self._make_processor() + + self.fd_config.model_config.think_token_sequences = { + "start": self.START_SEQUENCES, + "end": self.END_SEQUENCES, + "forced_end": [self.SEQ_VOCAB_SIZE], + } + with self.assertRaisesRegex(ValueError, "vocabulary range"): + self._make_processor() + + def test_sequence_prompt_scan_detects_contextual_start_variant(self): + req_id = "req_seq_prompt" + prompt_ids = [1, 12, 11, 5] + sampling_params = SamplingParams(logits_processors_args={"thinking_budget": 4}) + mock_req = MockRequest(req_id, prompt_ids, sampling_params) + + mock_runner = MockModelRunner(self.fd_config, max_num_seqs=1) + mock_runner.update_request_state(0, mock_req, pre_id=prompt_ids[-1]) + + processor = self._make_processor() + processor.update_state(mock_runner.share_inputs) + + self.assertTrue(processor._states[req_id].started) + self.assertFalse(processor._states[req_id].ended) + self.assertEqual(processor._states[req_id].tokens_after_start, 0) + + def test_sequence_natural_end_excludes_marker_tokens_from_budget(self): + req_id = "req_seq_natural_end" + prompt_ids = [10, 11] + sampling_params = SamplingParams(logits_processors_args={"thinking_budget": 8}) + mock_req = MockRequest(req_id, prompt_ids, sampling_params) + + mock_runner = MockModelRunner(self.fd_config, max_num_seqs=1) + mock_runner.update_request_state(0, mock_req, pre_id=prompt_ids[-1]) + + processor = self._make_processor() + processor.update_state(mock_runner.share_inputs) + self.assertTrue(processor._states[req_id].started) + + for token in (5, 20, 11): + mock_runner.update_request_state(0, mock_req, pre_id=token) + processor.update_state(mock_runner.share_inputs) + + state = processor._states[req_id] + self.assertTrue(state.ended) + # Only the payload token counts; the two end-marker tokens do not. + self.assertEqual(state.tokens_after_start, 1) + + logits = self._get_initial_logits(1) + processor.update_state(mock_runner.share_inputs) + processed_logits = processor.apply(logits) + self.assertEqual(paddle.argmax(processed_logits, axis=-1).item(), 0) + + def test_sequence_budget_forces_full_forced_end_sequence(self): + req_id = "req_seq_budget" + prompt_ids = [10, 11] + sampling_params = SamplingParams(logits_processors_args={"thinking_budget": 2}) + mock_req = MockRequest(req_id, prompt_ids, sampling_params) + + mock_runner = MockModelRunner(self.fd_config, max_num_seqs=1) + mock_runner.update_request_state(0, mock_req, pre_id=prompt_ids[-1]) + + processor = self._make_processor() + processor.update_state(mock_runner.share_inputs) + + for token in (5, 6): + mock_runner.update_request_state(0, mock_req, pre_id=token) + processor.update_state(mock_runner.share_inputs) + + forced_emissions = [] + for _ in self.FORCED_END_IDS: + logits = self._get_initial_logits(1) + processor.update_state(mock_runner.share_inputs) + processed_logits = processor.apply(logits) + next_token = mock_runner.generate_next_token(processed_logits)[0] + forced_emissions.append(next_token) + mock_runner.update_request_state(0, mock_req, pre_id=next_token) + + self.assertEqual(forced_emissions, self.FORCED_END_IDS) + + processor.update_state(mock_runner.share_inputs) + self.assertTrue(processor._states[req_id].ended) + logits = self._get_initial_logits(1) + processor.update_state(mock_runner.share_inputs) + processed_logits = processor.apply(logits) + self.assertEqual(paddle.argmax(processed_logits, axis=-1).item(), 0) + + def test_sequence_decode_time_start_detection(self): + req_id = "req_seq_decode_start" + prompt_ids = [1, 2, 3] + sampling_params = SamplingParams(logits_processors_args={"thinking_budget": 2}) + mock_req = MockRequest(req_id, prompt_ids, sampling_params) + + mock_runner = MockModelRunner(self.fd_config, max_num_seqs=1) + mock_runner.update_request_state(0, mock_req, pre_id=prompt_ids[-1], set_next_token=False) + + processor = self._make_processor() + processor.update_state(mock_runner.share_inputs) + self.assertFalse(processor._states[req_id].started) + + mock_runner.update_request_state(0, mock_req, pre_id=12) + processor.update_state(mock_runner.share_inputs) + self.assertFalse(processor._states[req_id].started) + + mock_runner.update_request_state(0, mock_req, pre_id=11) + processor.update_state(mock_runner.share_inputs) + self.assertTrue(processor._states[req_id].started) + self.assertEqual(processor._states[req_id].tokens_after_start, 0) + + def test_sequence_prompt_tail_prefix_is_not_consumed_twice(self): + self.fd_config.model_config.think_token_sequences = { + "start": [[10, 10]], + "end": self.END_SEQUENCES, + "forced_end": self.FORCED_END_IDS, + } + req_id = "req_seq_prompt_prefix" + sampling_params = SamplingParams(logits_processors_args={"thinking_budget": 2}) + mock_req = MockRequest(req_id, [10], sampling_params) + mock_runner = MockModelRunner(self.fd_config, max_num_seqs=1) + mock_runner.update_request_state(0, mock_req, pre_id=10) + processor = self._make_processor() + + processor.update_state(mock_runner.share_inputs) + + self.assertFalse(processor._states[req_id].started) + self.assertEqual(processor._states[req_id].start_buffer, [10]) + + mock_runner.update_request_state(0, mock_req, pre_id=10) + processor.update_state(mock_runner.share_inputs) + self.assertTrue(processor._states[req_id].started) + + def test_sequence_forcing_rejects_sampled_token_mismatch(self): + req_id = "req_seq_force_mismatch" + sampling_params = SamplingParams(logits_processors_args={"thinking_budget": 0}) + mock_req = MockRequest(req_id, [10, 11], sampling_params) + mock_runner = MockModelRunner(self.fd_config, max_num_seqs=1) + mock_runner.update_request_state(0, mock_req, pre_id=11) + processor = self._make_processor() + processor.update_state(mock_runner.share_inputs) + processor.apply(self._get_initial_logits(1)) + + mock_runner.update_request_state(0, mock_req, pre_id=9) + + with self.assertRaisesRegex(RuntimeError, "expected 30, received 9"): + processor.update_state(mock_runner.share_inputs) + + def test_sequence_stop_sentence_then_forced_end(self): + req_id = "req_seq_stop_sentence" + prompt_ids = [10, 11] + sampling_params = SamplingParams( + logits_processors_args={ + "thinking_budget": 3, + "think_stop_sentence_token_ids": [7, 8], + } + ) + mock_req = MockRequest(req_id, prompt_ids, sampling_params) + + mock_runner = MockModelRunner(self.fd_config, max_num_seqs=1) + mock_runner.update_request_state(0, mock_req, pre_id=prompt_ids[-1]) + + processor = self._make_processor() + processor.update_state(mock_runner.share_inputs) + + mock_runner.update_request_state(0, mock_req, pre_id=5) + processor.update_state(mock_runner.share_inputs) + + expected = [7, 8] + self.FORCED_END_IDS + for expected_token in expected: + logits = self._get_initial_logits(1) + processor.update_state(mock_runner.share_inputs) + processed_logits = processor.apply(logits) + next_token = mock_runner.generate_next_token(processed_logits)[0] + self.assertEqual(next_token, expected_token) + mock_runner.update_request_state(0, mock_req, pre_id=next_token) + + processor.update_state(mock_runner.share_inputs) + self.assertTrue(processor._states[req_id].ended) + + def test_sequence_new_round_resets_stop_sentence_progress(self): + req_id = "req_seq_stop_sentence_restart" + sampling_params = SamplingParams( + logits_processors_args={ + "thinking_budget": 1, + "think_stop_sentence_token_ids": [7, 8], + } + ) + mock_req = MockRequest(req_id, [10, 11], sampling_params) + mock_runner = MockModelRunner(self.fd_config, max_num_seqs=1) + mock_runner.update_request_state(0, mock_req, pre_id=11) + processor = self._make_processor() + processor.update_state(mock_runner.share_inputs) + + for expected_token in [7, 8] + self.FORCED_END_IDS: + logits = processor.apply(self._get_initial_logits(1)) + next_token = mock_runner.generate_next_token(logits)[0] + self.assertEqual(next_token, expected_token) + mock_runner.update_request_state(0, mock_req, pre_id=next_token) + processor.update_state(mock_runner.share_inputs) + + for token in (12, 11): + mock_runner.update_request_state(0, mock_req, pre_id=token) + processor.update_state(mock_runner.share_inputs) + + logits = processor.apply(self._get_initial_logits(1)) + self.assertEqual(mock_runner.generate_next_token(logits)[0], 7) + + class DummyTokenizerForTextProcessor: def __init__(self): self.vocab = {"x": 0} diff --git a/tests/operators/test_infllmv2_attention_forward.py b/tests/operators/test_infllmv2_attention_forward.py new file mode 100644 index 00000000000..fd78202058f --- /dev/null +++ b/tests/operators/test_infllmv2_attention_forward.py @@ -0,0 +1,731 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import pytest + +BLOCK_SIZE = 8 +KERNEL_SIZE = 4 +KERNEL_STRIDE = 2 +QUERY_HEADS = 4 +KV_HEADS = 2 +HEAD_DIM = 8 + + +def _load_ops(): + paddle = pytest.importorskip("paddle") + if not paddle.is_compiled_with_cuda(): + pytest.skip("InfLLM-V2 custom ops require CUDA Paddle.") + try: + from fastdeploy.model_executor.ops.gpu import ( + infllmv2_attention_forward, + infllmv2_select_blocks, + infllmv2_update_compressed_k, + ) + except ImportError: + pytest.skip("InfLLM-V2 custom ops are not present in the installed fastdeploy_ops package.") + paddle.set_device("gpu") + return paddle, infllmv2_update_compressed_k, infllmv2_select_blocks, infllmv2_attention_forward + + +def _unpack(outputs): + return outputs if isinstance(outputs, (tuple, list)) else (outputs,) + + +def _make_paged_cache(rng, batch_size=2, blocks_per_sequence=5): + physical_blocks = batch_size * blocks_per_sequence + block_tables = np.arange(physical_blocks, dtype=np.int32).reshape(batch_size, blocks_per_sequence) + block_tables[0] = block_tables[0, [3, 0, 4, 1, 2]] + block_tables[1] = block_tables[1, [2, 4, 0, 3, 1]] + logical_k = rng.normal(size=(batch_size, blocks_per_sequence * BLOCK_SIZE, KV_HEADS, HEAD_DIM)).astype("float32") + logical_v = rng.normal(size=logical_k.shape).astype("float32") + key_cache = np.zeros((physical_blocks, KV_HEADS, BLOCK_SIZE, HEAD_DIM), dtype="float32") + value_cache = np.zeros_like(key_cache) + for batch_id in range(batch_size): + for logical_block, physical_block in enumerate(block_tables[batch_id]): + begin = logical_block * BLOCK_SIZE + end = begin + BLOCK_SIZE + key_cache[physical_block] = logical_k[batch_id, begin:end].transpose(1, 0, 2) + value_cache[physical_block] = logical_v[batch_id, begin:end].transpose(1, 0, 2) + return logical_k, logical_v, key_cache, value_cache, block_tables + + +def _metadata(paddle, batch_size, blocks_per_sequence, position): + return ( + paddle.full([batch_size], position, dtype="int32"), + paddle.ones([batch_size], dtype="int32"), + paddle.arange(batch_size, dtype="int32"), + paddle.arange(batch_size + 1, dtype="int32"), + ) + + +def _build_summaries(paddle, update, key_cache, block_tables, sequence_length, return_workspaces=False): + batch_size = block_tables.shape[0] + physical_blocks = key_cache.shape[0] + tokens = batch_size * sequence_length + current = paddle.zeros([tokens, QUERY_HEADS, HEAD_DIM], dtype=key_cache.dtype) + seq_decoder = paddle.zeros([batch_size], dtype="int32") + seq_now = paddle.full([batch_size], sequence_length, dtype="int32") + batch_ids = paddle.repeat_interleave(paddle.arange(batch_size, dtype="int32"), sequence_length) + cu = paddle.arange(batch_size + 1, dtype="int32") * sequence_length + fine = paddle.zeros([physical_blocks, KV_HEADS, BLOCK_SIZE // KERNEL_STRIDE, HEAD_DIM], dtype=key_cache.dtype) + coarse = paddle.zeros( + [physical_blocks, KV_HEADS, BLOCK_SIZE // (4 * KERNEL_STRIDE), HEAD_DIM], dtype=key_cache.dtype + ) + outputs = _unpack( + update( + current, + key_cache, + fine, + coarse, + block_tables, + seq_decoder, + seq_now, + batch_ids, + cu, + KERNEL_SIZE, + KERNEL_STRIDE, + ) + ) + if return_workspaces: + return outputs[0], outputs[1], fine, coarse + return outputs[0], outputs[1] + + +def _selection_workspaces(paddle, tokens, blocks_per_sequence, capacity, dtype="float32"): + max_coarse_windows = max( + 0, + (blocks_per_sequence * BLOCK_SIZE - 4 * KERNEL_SIZE) // (4 * KERNEL_STRIDE) + 1, + ) + coarse_splits = max(1, (max_coarse_windows + 15) // 16) + return ( + paddle.empty([tokens, KV_HEADS, capacity], dtype="int32"), + paddle.empty([tokens, KV_HEADS, blocks_per_sequence], dtype="float32"), + paddle.empty([tokens, KV_HEADS], dtype="int32"), + paddle.empty([tokens, QUERY_HEADS], dtype="float32"), + paddle.empty([tokens, QUERY_HEADS, coarse_splits], dtype="float32"), + paddle.empty([tokens, QUERY_HEADS, coarse_splits], dtype="float32"), + ) + + +def _attention_workspaces(paddle, tokens, capacity, dtype): + splits = (capacity + 1) // 2 + return ( + paddle.empty([tokens, QUERY_HEADS, HEAD_DIM], dtype=dtype), + paddle.empty([tokens, QUERY_HEADS, splits, HEAD_DIM], dtype="float32"), + paddle.empty([tokens, QUERY_HEADS, splits], dtype="float32"), + paddle.empty([tokens, QUERY_HEADS, splits], dtype="float32"), + ) + + +def _reference_attention(query, logical_k, logical_v, selected, positions): + output = np.zeros_like(query, dtype="float32") + group_size = QUERY_HEADS // KV_HEADS + for token_id in range(query.shape[0]): + batch_id = token_id + for query_head in range(QUERY_HEADS): + kv_head = query_head // group_size + blocks = selected[token_id, kv_head] + blocks = blocks[blocks >= 0] + indices = np.concatenate([np.arange(block * BLOCK_SIZE, (block + 1) * BLOCK_SIZE) for block in blocks]) + indices = indices[indices <= positions[token_id]] + if not len(indices): + continue + logits = logical_k[batch_id, indices, kv_head] @ query[token_id, query_head] + logits = logits / np.sqrt(HEAD_DIM) + probabilities = np.exp(logits - logits.max()) + probabilities /= probabilities.sum() + output[token_id, query_head] = probabilities @ logical_v[batch_id, indices, kv_head] + return output + + +def _reference_stage1(query, logical_k, position, topk, init_blocks, local_blocks): + tokens = query.shape[0] + blocks_per_sequence = logical_k.shape[1] // BLOCK_SIZE + group_size = QUERY_HEADS // KV_HEADS + scale = 1.0 / np.sqrt(HEAD_DIM) + visible_length = position + 1 + + fine = np.stack( + [ + logical_k[:, start : start + KERNEL_SIZE].mean(axis=1) + for start in range(0, visible_length - KERNEL_SIZE + 1, KERNEL_STRIDE) + ], + axis=1, + ) + coarse_kernel = 4 * KERNEL_SIZE + coarse_stride = 4 * KERNEL_STRIDE + coarse = np.stack( + [ + logical_k[:, start : start + coarse_kernel].mean(axis=1) + for start in range(0, visible_length - coarse_kernel + 1, coarse_stride) + ], + axis=1, + ) + + coarse_lse = np.empty((tokens, QUERY_HEADS), dtype=np.float32) + block_scores = np.full((tokens, KV_HEADS, blocks_per_sequence), -np.inf, dtype=np.float32) + selected = np.full((tokens, KV_HEADS, topk + local_blocks), -1, dtype=np.int32) + current_block = position // BLOCK_SIZE + fine_slots = BLOCK_SIZE // KERNEL_STRIDE + for token_id in range(tokens): + batch_id = token_id + for query_head in range(QUERY_HEADS): + kv_head = query_head // group_size + logits = coarse[batch_id, :, kv_head] @ query[token_id, query_head] * scale + coarse_lse[token_id, query_head] = np.logaddexp.reduce(logits) + + for kv_head in range(KV_HEADS): + for logical_block in range(blocks_per_sequence): + if logical_block < init_blocks or ( + logical_block <= current_block and logical_block + local_blocks > current_block + ): + block_scores[token_id, kv_head, logical_block] = np.inf + continue + first_window = max(0, logical_block * fine_slots - 1) + last_window = min(fine.shape[1], (logical_block + 1) * fine_slots) + window_scores = [] + for window in range(first_window, last_window): + gqa_score = 0.0 + for group_head in range(group_size): + query_head = kv_head * group_size + group_head + logit = fine[batch_id, window, kv_head] @ query[token_id, query_head] * scale + gqa_score += np.exp(logit - coarse_lse[token_id, query_head]) + window_scores.append(gqa_score) + block_scores[token_id, kv_head, logical_block] = max(window_scores, default=-np.inf) + + ranked = sorted( + range(blocks_per_sequence), + key=lambda block: (-block_scores[token_id, kv_head, block], block), + )[: topk + local_blocks] + selected[token_id, kv_head, : len(ranked)] = sorted(ranked) + return block_scores, coarse_lse, selected + + +def test_infllmv2_update_compressed_k_handles_cross_page_noncontiguous_tables(): + paddle, update, _, _ = _load_ops() + rng = np.random.default_rng(2026) + logical_k, _, key_np, _, table_np = _make_paged_cache(rng) + key = paddle.to_tensor(key_np) + tables = paddle.to_tensor(table_np) + + fine, coarse, fine_workspace, coarse_workspace = _build_summaries( + paddle, update, key, tables, 40, return_workspaces=True + ) + fine_np, coarse_np = fine.numpy(), coarse.numpy() + + # Fine [6, 9] and coarse [0, 15] cross a logical page boundary and + # are stored in the physical page containing the final token. + physical_page = table_np[0, 1] + np.testing.assert_allclose(fine_np[physical_page, :, 0], logical_k[0, 6:10].mean(axis=0), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(coarse_np[physical_page, :, 0], logical_k[0, 0:16].mean(axis=0), rtol=1e-5, atol=1e-5) + assert fine._is_shared_buffer_with(fine_workspace) + assert coarse._is_shared_buffer_with(coarse_workspace) + + +def test_infllmv2_stage1_selects_per_request_and_kv_head_with_sorted_padding(): + paddle, update, select, _ = _load_ops() + rng = np.random.default_rng(17) + _, _, key_np, _, table_np = _make_paged_cache(rng) + key = paddle.to_tensor(key_np) + tables = paddle.to_tensor(table_np) + fine, coarse = _build_summaries(paddle, update, key, tables, 40) + query = paddle.to_tensor(rng.normal(size=(2, QUERY_HEADS, HEAD_DIM)).astype("float32")) + seq_decoder, seq_now, batch_ids, cu = _metadata(paddle, 2, 5, 39) + workspaces = _selection_workspaces(paddle, 2, 5, 3) + + outputs = _unpack( + select( + query, + fine, + coarse, + tables, + seq_decoder, + seq_now, + batch_ids, + cu, + *workspaces, + BLOCK_SIZE, + KERNEL_SIZE, + KERNEL_STRIDE, + 2, + 16, + 1, + 1, + ) + ) + selected, block_scores, counts = outputs[0].numpy(), outputs[1].numpy(), outputs[2].numpy() + + assert np.all(counts == 3) + # local_blocks includes the current block. With local_blocks=1, only block 4 + # is forced local; block 3 must retain its finite semantic score. + assert np.all(np.isposinf(block_scores[:, :, 0])) + assert np.all(np.isfinite(block_scores[:, :, 3])) + assert np.all(np.isposinf(block_scores[:, :, 4])) + assert np.all(selected[:, :, 0] == 0) + assert np.all(selected[:, :, -1] == 4) + for returned, workspace in zip(outputs, workspaces): + assert returned._is_shared_buffer_with(workspace) + + +def test_infllmv2_stage1_short_context_selects_all_visible_blocks(): + paddle, update, select, _ = _load_ops() + rng = np.random.default_rng(23) + _, _, key_np, _, table_np = _make_paged_cache(rng) + key = paddle.to_tensor(key_np) + tables = paddle.to_tensor(table_np) + fine, coarse = _build_summaries(paddle, update, key, tables, 40) + query = paddle.to_tensor(rng.normal(size=(2, QUERY_HEADS, HEAD_DIM)).astype("float32")) + seq_decoder, seq_now, batch_ids, cu = _metadata(paddle, 2, 5, 15) + workspaces = _selection_workspaces(paddle, 2, 5, 4) + + outputs = _unpack( + select( + query, + fine, + coarse, + tables, + seq_decoder, + seq_now, + batch_ids, + cu, + *workspaces, + BLOCK_SIZE, + KERNEL_SIZE, + KERNEL_STRIDE, + 2, + 32, + 1, + 1, + ) + ) + + np.testing.assert_array_equal(outputs[0].numpy(), np.array([[[0, 1, -1, -1], [0, 1, -1, -1]]] * 2)) + np.testing.assert_array_equal(outputs[2].numpy(), np.full([2, 2], 2)) + + +def test_infllmv2_stage1_scores_and_topk_match_numpy_reference(): + paddle, update, select, _ = _load_ops() + rng = np.random.default_rng(29) + logical_k, _, key_np, _, table_np = _make_paged_cache(rng) + key = paddle.to_tensor(key_np) + tables = paddle.to_tensor(table_np) + fine, coarse = _build_summaries(paddle, update, key, tables, 40) + query_np = rng.normal(size=(2, QUERY_HEADS, HEAD_DIM)).astype("float32") + query = paddle.to_tensor(query_np) + seq_decoder, seq_now, batch_ids, cu = _metadata(paddle, 2, 5, 39) + workspaces = _selection_workspaces(paddle, 2, 5, 2) + + outputs = _unpack( + select( + query, + fine, + coarse, + tables, + seq_decoder, + seq_now, + batch_ids, + cu, + *workspaces, + BLOCK_SIZE, + KERNEL_SIZE, + KERNEL_STRIDE, + 2, + 16, + 0, + 0, + ) + ) + expected_scores, expected_lse, expected_selected = _reference_stage1(query_np, logical_k, 39, 2, 0, 0) + + np.testing.assert_allclose(outputs[1].numpy(), expected_scores, rtol=2e-5, atol=2e-5) + np.testing.assert_allclose(outputs[3].numpy(), expected_lse, rtol=2e-5, atol=2e-5) + np.testing.assert_array_equal(outputs[0].numpy(), expected_selected) + np.testing.assert_array_equal(outputs[2].numpy(), np.full([2, KV_HEADS], 2)) + + +def test_infllmv2_tensor_core_scores_topk_and_attention_match_numpy_reference(): + paddle, update, select, attention = _load_ops() + rng = np.random.default_rng(53) + block_size = 64 + kernel_size = 32 + kernel_stride = 16 + query_heads = 32 + kv_heads = 2 + head_dim = 128 + blocks = 4 + sequence_length = blocks * block_size + block_table_np = np.array([[2, 0, 3, 1]], dtype=np.int32) + logical_k = (0.25 * rng.normal(size=(sequence_length, kv_heads, head_dim))).astype("float32") + key_np = np.zeros([blocks, kv_heads, block_size, head_dim], dtype="float32") + for logical_block, physical_block in enumerate(block_table_np[0]): + begin = logical_block * block_size + key_np[physical_block] = logical_k[begin : begin + block_size].transpose(1, 0, 2) + + key = paddle.to_tensor(key_np).astype("bfloat16") + block_tables = paddle.to_tensor(block_table_np) + current = paddle.zeros([sequence_length, query_heads, head_dim], dtype="bfloat16") + fine = paddle.zeros([blocks, kv_heads, block_size // kernel_stride, head_dim], dtype="bfloat16") + coarse = paddle.zeros([blocks, kv_heads, block_size // (4 * kernel_stride), head_dim], dtype="bfloat16") + update_outputs = _unpack( + update( + current, + key, + fine, + coarse, + block_tables, + paddle.zeros([1], dtype="int32"), + paddle.full([1], sequence_length, dtype="int32"), + paddle.zeros([sequence_length], dtype="int32"), + paddle.to_tensor([0, sequence_length], dtype="int32"), + kernel_size, + kernel_stride, + ) + ) + fine_np = update_outputs[0].astype("float32").numpy() + coarse_np = update_outputs[1].astype("float32").numpy() + + query = paddle.to_tensor((0.25 * rng.normal(size=(1, query_heads, head_dim))).astype("float32")).astype("bfloat16") + query_np = query.astype("float32").numpy()[0] + selected = paddle.empty([1, kv_heads, 2], dtype="int32") + block_scores = paddle.empty([1, kv_heads, blocks], dtype="float32") + selected_counts = paddle.empty([1, kv_heads], dtype="int32") + coarse_lse = paddle.empty([1, query_heads], dtype="float32") + coarse_partial_max = paddle.empty([1, query_heads, 1], dtype="float32") + coarse_partial_sum = paddle.empty([1, query_heads, 1], dtype="float32") + outputs = _unpack( + select( + query, + update_outputs[0], + update_outputs[1], + block_tables, + paddle.full([1], sequence_length - 1, dtype="int32"), + paddle.ones([1], dtype="int32"), + paddle.zeros([1], dtype="int32"), + paddle.to_tensor([0, 1], dtype="int32"), + selected, + block_scores, + selected_counts, + coarse_lse, + coarse_partial_max, + coarse_partial_sum, + block_size, + kernel_size, + kernel_stride, + 2, + 0, + 0, + 0, + ) + ) + + fine_windows = [] + for window_end in range(kernel_size - 1, sequence_length, kernel_stride): + logical_block = window_end // block_size + physical_block = block_table_np[0, logical_block] + slot = (window_end % block_size) // kernel_stride + fine_windows.append(fine_np[physical_block, :, slot]) + coarse_windows = [] + for window_end in range(4 * kernel_size - 1, sequence_length, 4 * kernel_stride): + logical_block = window_end // block_size + physical_block = block_table_np[0, logical_block] + slot = (window_end % block_size) // (4 * kernel_stride) + coarse_windows.append(coarse_np[physical_block, :, slot]) + fine_windows = np.stack(fine_windows) + coarse_windows = np.stack(coarse_windows) + scale = 1.0 / np.sqrt(head_dim) + expected_lse = np.empty([query_heads], dtype="float32") + for query_head in range(query_heads): + kv_head = query_head // (query_heads // kv_heads) + logits = coarse_windows[:, kv_head] @ query_np[query_head] * scale + expected_lse[query_head] = np.logaddexp.reduce(logits) + expected_scores = np.empty([kv_heads, blocks], dtype="float32") + fine_slots = block_size // kernel_stride + for kv_head in range(kv_heads): + for logical_block in range(blocks): + first_window = max(0, logical_block * fine_slots - 1) + last_window = min(len(fine_windows), (logical_block + 1) * fine_slots) + window_scores = [] + for window in range(first_window, last_window): + score = 0.0 + for group_head in range(query_heads // kv_heads): + query_head = kv_head * (query_heads // kv_heads) + group_head + logit = fine_windows[window, kv_head] @ query_np[query_head] * scale + score += np.exp(logit - expected_lse[query_head]) + window_scores.append(score) + expected_scores[kv_head, logical_block] = max(window_scores) + expected_selected = np.empty([kv_heads, 2], dtype="int32") + for kv_head in range(kv_heads): + ranked = sorted(range(blocks), key=lambda block: (-expected_scores[kv_head, block], block))[:2] + expected_selected[kv_head] = sorted(ranked) + + np.testing.assert_allclose(outputs[1].numpy()[0], expected_scores, rtol=2e-2, atol=2e-3) + np.testing.assert_allclose(outputs[3].numpy()[0], expected_lse, rtol=2e-2, atol=2e-3) + np.testing.assert_array_equal(outputs[0].numpy()[0], expected_selected) + np.testing.assert_array_equal(outputs[2].numpy(), np.full([1, kv_heads], 2)) + + attention_output = paddle.empty([1, query_heads, head_dim], dtype="bfloat16") + partial_acc = paddle.empty([1, query_heads, 1, head_dim], dtype="float32") + partial_max = paddle.empty([1, query_heads, 1], dtype="float32") + partial_sum = paddle.empty([1, query_heads, 1], dtype="float32") + attention_outputs = _unpack( + attention( + query, + key, + key, + block_tables, + paddle.full([1], sequence_length - 1, dtype="int32"), + paddle.ones([1], dtype="int32"), + paddle.zeros([1], dtype="int32"), + paddle.to_tensor([0, 1], dtype="int32"), + outputs[0], + attention_output, + partial_acc, + partial_max, + partial_sum, + ) + ) + key_quantized = key.astype("float32").numpy() + logical_quantized = np.empty([sequence_length, kv_heads, head_dim], dtype="float32") + for logical_block, physical_block in enumerate(block_table_np[0]): + begin = logical_block * block_size + logical_quantized[begin : begin + block_size] = key_quantized[physical_block].transpose(1, 0, 2) + expected_attention = np.empty([query_heads, head_dim], dtype="float32") + selected_np = outputs[0].numpy()[0] + for query_head in range(query_heads): + kv_head = query_head // (query_heads // kv_heads) + indices = np.concatenate( + [ + np.arange(logical_block * block_size, (logical_block + 1) * block_size) + for logical_block in selected_np[kv_head] + ] + ) + logits = logical_quantized[indices, kv_head] @ query_np[query_head] * scale + probabilities = np.exp(logits - logits.max()) + probabilities /= probabilities.sum() + expected_attention[query_head] = probabilities @ logical_quantized[indices, kv_head] + np.testing.assert_allclose( + attention_outputs[0].astype("float32").numpy()[0], + expected_attention, + rtol=4e-2, + atol=4e-2, + ) + + +@pytest.mark.parametrize("query_tile_size", [63, 127, 128]) +def test_infllmv2_sparse_prefill_tile_matches_dense_causal_attention(query_tile_size): + paddle, _, _, _ = _load_ops() + from paddle.nn.functional.flash_attention import flash_attn_unpadded + + from fastdeploy.model_executor.layers.attention.infllmv2_attention_backend import ( + InfLLMV2AttentionBackend, + ) + + paddle.seed(2026) + sparse_start = 128 + sequence_length = sparse_start + query_tile_size + block_size = 64 + query_heads = 32 + kv_heads = 2 + head_dim = 128 + backend = InfLLMV2AttentionBackend.__new__(InfLLMV2AttentionBackend) + backend.block_size = block_size + backend.num_heads = query_heads + backend.kv_num_heads = kv_heads + backend.head_dim = head_dim + blocks = (sequence_length + block_size - 1) // block_size + backend.topk = blocks + backend.local_blocks = 0 + + query = paddle.randn([sequence_length, query_heads, head_dim], dtype="bfloat16") + key = paddle.randn([sequence_length, kv_heads, head_dim], dtype="bfloat16") + value = paddle.randn(key.shape, dtype="bfloat16") + physical_order = [2, 0, 1] if blocks == 3 else [2, 0, 3, 1] + block_table = paddle.to_tensor([physical_order], dtype="int32") + key_cache = paddle.zeros([blocks, kv_heads, block_size, head_dim], dtype="bfloat16") + value_cache = paddle.zeros_like(key_cache) + for logical_block, physical_block in enumerate(block_table.numpy()[0]): + begin = logical_block * block_size + end = min(sequence_length, begin + block_size) + key_cache[physical_block, :, : end - begin] = paddle.transpose(key[begin:end], [1, 0, 2]) + value_cache[physical_block, :, : end - begin] = paddle.transpose(value[begin:end], [1, 0, 2]) + + selected = paddle.tile( + paddle.arange(blocks, dtype="int32").reshape([1, 1, blocks]), + [1, kv_heads, 1], + ) + output = backend._sparse_prefill_tile_batch( + query[sparse_start:], + key_cache, + value_cache, + block_table, + selected, + sparse_start, + query_tile_size, + ) + cu_seqlens = paddle.to_tensor([0, sequence_length], dtype="int32") + dense = flash_attn_unpadded( + query, + key, + value, + cu_seqlens, + cu_seqlens, + sequence_length, + sequence_length, + scale=head_dim**-0.5, + causal=True, + training=False, + )[0] + + np.testing.assert_allclose( + output.astype("float32").numpy(), + dense[sparse_start:].astype("float32").numpy(), + rtol=4e-2, + atol=4e-2, + ) + + +@pytest.mark.parametrize(("dtype", "rtol", "atol"), [("float32", 2e-5, 2e-5), ("bfloat16", 4e-2, 4e-2)]) +def test_infllmv2_stage2_maps_logical_blocks_to_paged_cache(dtype, rtol, atol): + paddle, _, _, attention = _load_ops() + rng = np.random.default_rng(31) + logical_k, logical_v, key_np, value_np, table_np = _make_paged_cache(rng) + query_np = rng.normal(size=(2, QUERY_HEADS, HEAD_DIM)).astype("float32") + selected_np = np.array([[[0, 2, 4], [0, 2, 4]], [[1, 3, 4], [1, 3, 4]]], np.int32) + positions = np.array([39, 39], np.int32) + query = paddle.to_tensor(query_np).astype(dtype) + key = paddle.to_tensor(key_np).astype(dtype) + value = paddle.to_tensor(value_np).astype(dtype) + tables = paddle.to_tensor(table_np) + seq_decoder, seq_now, batch_ids, cu = _metadata(paddle, 2, 5, 39) + workspaces = _attention_workspaces(paddle, 2, 3, dtype) + + outputs = _unpack( + attention( + query, + key, + value, + tables, + seq_decoder, + seq_now, + batch_ids, + cu, + paddle.to_tensor(selected_np), + *workspaces, + ) + ) + expected = _reference_attention(query_np, logical_k, logical_v, selected_np, positions) + + np.testing.assert_allclose(outputs[0].astype("float32").numpy(), expected, rtol=rtol, atol=atol) + for returned, workspace in zip(outputs, workspaces): + assert returned._is_shared_buffer_with(workspace) + + +def test_infllmv2_stage1_stage2_closed_loop_matches_dense_oracle_below_threshold(): + paddle, update, select, attention = _load_ops() + rng = np.random.default_rng(47) + logical_k, logical_v, key_np, value_np, table_np = _make_paged_cache(rng, blocks_per_sequence=5) + key = paddle.to_tensor(key_np) + value = paddle.to_tensor(value_np) + tables = paddle.to_tensor(table_np) + fine, coarse = _build_summaries(paddle, update, key, tables, 40) + query_np = rng.normal(size=(2, QUERY_HEADS, HEAD_DIM)).astype("float32") + query = paddle.to_tensor(query_np) + seq_decoder, seq_now, batch_ids, cu = _metadata(paddle, 2, 5, 15) + selection_ws = _selection_workspaces(paddle, 2, 5, 4) + selected = _unpack( + select( + query, + fine, + coarse, + tables, + seq_decoder, + seq_now, + batch_ids, + cu, + *selection_ws, + BLOCK_SIZE, + KERNEL_SIZE, + KERNEL_STRIDE, + 2, + 32, + 1, + 1, + ) + )[0] + attention_ws = _attention_workspaces(paddle, 2, 4, "float32") + output = _unpack( + attention( + query, + key, + value, + tables, + seq_decoder, + seq_now, + batch_ids, + cu, + selected, + *attention_ws, + ) + )[0] + selected_np = selected.numpy() + expected = _reference_attention(query_np, logical_k, logical_v, selected_np, np.array([15, 15])) + + np.testing.assert_allclose(output.numpy(), expected, rtol=2e-5, atol=2e-5) + + +def test_infllmv2_ops_reject_invalid_metadata_and_workspace_contracts(): + paddle, _, select, attention = _load_ops() + query = paddle.zeros([1, QUERY_HEADS, HEAD_DIM]) + compressed = paddle.zeros([2, KV_HEADS, BLOCK_SIZE // KERNEL_STRIDE, HEAD_DIM]) + compressed2 = paddle.zeros([2, KV_HEADS, BLOCK_SIZE // (4 * KERNEL_STRIDE), HEAD_DIM]) + tables = paddle.arange(2, dtype="int32").reshape([1, 2]) + seq_decoder, seq_now, batch_ids, cu = _metadata(paddle, 1, 2, 7) + selection_ws = _selection_workspaces(paddle, 1, 2, 2) + + with pytest.raises(Exception, match="int32 metadata"): + select( + query, + compressed, + compressed2, + tables.astype("int64"), + seq_decoder, + seq_now, + batch_ids, + cu, + *selection_ws, + BLOCK_SIZE, + KERNEL_SIZE, + KERNEL_STRIDE, + 1, + 16, + 0, + 0, + ) + + key = paddle.zeros([2, KV_HEADS, BLOCK_SIZE, HEAD_DIM]) + attention_ws = _attention_workspaces(paddle, 1, 2, "float32") + with pytest.raises(Exception, match="dtype int32"): + attention( + query, + key, + key, + tables, + seq_decoder, + seq_now, + batch_ids, + cu, + paddle.zeros([1, KV_HEADS, 2], dtype="int64"), + *attention_ws, + ) diff --git a/tests/quantization/test_minicpm41_int_quant.py b/tests/quantization/test_minicpm41_int_quant.py new file mode 100644 index 00000000000..d0808a31ddf --- /dev/null +++ b/tests/quantization/test_minicpm41_int_quant.py @@ -0,0 +1,117 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def read_repo_file(path): + return (REPO_ROOT / path).read_text(encoding="utf-8") + + +def load_module(module_path): + spec = importlib.util.spec_from_file_location(module_path.stem, module_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def make_parse_inputs(quantization): + args = SimpleNamespace( + quantization=quantization, + dynamic_load_weight=False, + enable_mega_moe=False, + ) + model_config = SimpleNamespace( + architectures=["MiniCPMForCausalLM"], + quantization_config=None, + model_format="torch", + is_quantized=False, + ) + return args, model_config + + +def test_minicpm41_declares_only_verified_wint_quantizations(): + module = load_module(REPO_ROOT / "fastdeploy/model_executor/models/minicpm41/config_minicpm41.py") + + assert module.SUPPORTED_QUANTIZATIONS == ("wint4", "wint8") + + +@pytest.mark.parametrize( + ("quant_name", "expected_class", "expected_algo"), + [ + ("wint4", "WINT4Config", "weight_only_int4"), + ("wint8", "WINT8Config", "weight_only_int8"), + ], +) +def test_minicpm41_cli_wint_uses_bf16_checkpoint_online_quantization( + monkeypatch, + quant_name, + expected_class, + expected_algo, +): + pytest.importorskip("paddle") + from fastdeploy.model_executor.layers.quantization import parse_quant_config + + class FakeQuantConfig: + @classmethod + def from_config(cls, config): + return SimpleNamespace( + class_name=expected_class, + quant_name=quant_name, + algo=expected_algo, + is_checkpoint_bf16=not config.get("is_quantized", False), + ) + + monkeypatch.setattr( + "fastdeploy.model_executor.layers.quantization.get_quantization_config", + lambda name: FakeQuantConfig, + ) + + args, model_config = make_parse_inputs(quant_name) + quant_config = parse_quant_config(args, model_config, is_ernie=False, is_v1_loader=False) + + assert quant_config.class_name == expected_class + assert quant_config.quant_name == quant_name + assert quant_config.algo == expected_algo + assert quant_config.is_checkpoint_bf16 is True + assert model_config.is_quantized is False + assert args.quantization == {"quantization": quant_name} + + +@pytest.mark.parametrize("quant_name", ["wint4", "wint8"]) +def test_minicpm41_one_key_quant_dict_stays_online(monkeypatch, quant_name): + pytest.importorskip("paddle") + from fastdeploy.model_executor.layers.quantization import parse_quant_config + + class FakeQuantConfig: + @classmethod + def from_config(cls, config): + return SimpleNamespace(is_checkpoint_bf16=not config.get("is_quantized", False)) + + monkeypatch.setattr( + "fastdeploy.model_executor.layers.quantization.get_quantization_config", + lambda name: FakeQuantConfig, + ) + + args, model_config = make_parse_inputs({"quantization": quant_name}) + quant_config = parse_quant_config(args, model_config, is_ernie=False, is_v1_loader=False) + + assert quant_config.is_checkpoint_bf16 is True + assert model_config.is_quantized is False diff --git a/tests/quantization/test_minicpm41_quality_eval.py b/tests/quantization/test_minicpm41_quality_eval.py new file mode 100644 index 00000000000..f34a0ed92da --- /dev/null +++ b/tests/quantization/test_minicpm41_quality_eval.py @@ -0,0 +1,64 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def load_quality_module(): + module_path = REPO_ROOT / "scripts/evaluate_minicpm41_quant_quality.py" + spec = importlib.util.spec_from_file_location("evaluate_minicpm41_quant_quality", module_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_score_output_accepts_nfkc_exact_answer(): + module = load_quality_module() + + passed, _ = module.score_output({"exact_answers": ("H2O",)}, " H₂O。 ") + + assert passed + + +def test_score_output_rejects_explanation_for_exact_answer(): + module = load_quality_module() + + passed, detail = module.score_output({"exact_answers": ("391",)}, "答案是 391。") + + assert not passed + assert "391" in detail + + +def test_score_output_requires_every_keyword(): + module = load_quality_module() + + passed, detail = module.score_output({"required_terms": ("蓝", "散射")}, "天空看起来是蓝色的。") + + assert not passed + assert "散射" in detail + + +def test_generation_config_records_deterministic_request_settings(): + module = load_quality_module() + + assert module.generation_config(96) == { + "max_tokens": 96, + "temperature": 0, + "top_p": 1, + "stream": False, + "chat_template_kwargs": {"enable_thinking": False}, + }