diff --git a/benchmarks/benchmark_rwkv7_quantization.py b/benchmarks/benchmark_rwkv7_quantization.py new file mode 100644 index 000000000000..961a9de807de --- /dev/null +++ b/benchmarks/benchmark_rwkv7_quantization.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark RWKV7 16-bit, online INT8, TorchAO INT8/INT4, and BitsAndBytes. + +Every setting runs in a fresh process. Besides median output throughput, the +report extracts vLLM's model-resident GPU-memory measurement and compares +greedy tokens with the FP16 run. The default gates encode the production +target: quantized model memory must decrease and throughput must not regress. +""" + +import argparse +import json +import statistics +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +import regex as re + +DEFAULT_PROMPTS = [ + "The Eiffel Tower is located in", + "A short proof that there are infinitely many primes begins", + "Write a Python function that computes Fibonacci numbers:", + "The most important property of recurrent neural networks is", + "Once upon a time in a quiet village", + "Explain why the sky appears blue during the day.", + "In numerical analysis, floating point accumulation order", + "Translate to Chinese: artificial intelligence inference engine", +] + +MODEL_MEMORY_PATTERN = re.compile(r"Model loading took ([0-9.]+) GiB memory") +NOISY_BNB_MESSAGES = ("MatMul8bitLt: inputs will be cast",) +SETTINGS = ( + "fp16", + "online-int8", + "torchao-int8", + "torchao-int4", + "bnb-int8", + "bnb-nf4", +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="FP16 reference checkpoint") + parser.add_argument("--tokenizer") + parser.add_argument("--int8-model", help="Pre-quantized BitsAndBytes INT8 model") + parser.add_argument( + "--int4-model", + help="Pre-quantized BitsAndBytes NF4 model; omit for inflight NF4", + ) + parser.add_argument( + "--settings", + nargs="+", + choices=SETTINGS, + default=("fp16", "online-int8"), + ) + parser.add_argument("--prompt", action="append", default=[]) + parser.add_argument("--prompts-file", type=Path) + parser.add_argument("--max-tokens", type=int, default=64) + parser.add_argument("--logprobs", type=int, default=1) + parser.add_argument( + "--warmup-runs", + type=int, + default=3, + help="full-engine warmups; three also settles lazy BitsAndBytes state", + ) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--dtype", default="half") + parser.add_argument("--max-model-len", type=int, default=2048) + parser.add_argument("--max-num-batched-tokens", type=int, default=32) + parser.add_argument("--gpu-memory-utilization", type=float, default=0.8) + parser.add_argument( + "--enforce-eager", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--async-scheduling", action=argparse.BooleanOptionalAction, default=False + ) + parser.add_argument( + "--ignore-eos", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--require-gates", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--require-repeatable", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument("--min-speed-ratio", type=float, default=1.0) + parser.add_argument("--min-memory-reduction", type=float, default=0.01) + parser.add_argument( + "--worker-setting", + choices=SETTINGS, + help=argparse.SUPPRESS, + ) + parser.add_argument("--worker-model", help=argparse.SUPPRESS) + parser.add_argument( + "--worker-inflight-nf4", action="store_true", help=argparse.SUPPRESS + ) + return parser.parse_args() + + +def load_prompts(args: argparse.Namespace) -> list[str]: + if args.warmup_runs < 0 or args.repeats < 1: + raise ValueError("--warmup-runs must be >= 0 and --repeats must be >= 1") + prompts = list(args.prompt) + if args.prompts_file is not None: + loaded = json.loads(args.prompts_file.read_text()) + if not isinstance(loaded, list) or not all(isinstance(x, str) for x in loaded): + raise ValueError("--prompts-file must contain a JSON array of strings") + prompts.extend(loaded) + return prompts or DEFAULT_PROMPTS + + +def run_worker(args: argparse.Namespace) -> None: + assert args.worker_setting is not None + assert args.worker_model is not None + + # Quantization benchmarks measure weight formats, not the experimental + # recurrent kernel. Keep the recurrent path deterministic and fail-closed. + import os + + os.environ["VLLM_RWKV7_KERNEL"] = "torch" + if args.require_repeatable: + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + # Keep the benchmark self-contained on CUDA hosts without a full nvcc + # toolchain; sampling is outside the measured model/kernel scope. + os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0") + + from vllm import LLM, SamplingParams + + if args.require_repeatable: + import torch + + torch.use_deterministic_algorithms(True) + + llm_kwargs: dict[str, Any] = {} + if args.worker_setting == "online-int8": + # vLLM's online-quant configuration was generalized from the original + # scheme/override API to QuantSpec. Supporting both keeps this benchmark + # usable while bisecting releases and downstream integration branches. + from vllm.config import quantization as online_quant_config + + if hasattr(online_quant_config, "QuantizationConfigArgs"): + quantization_config = {"linear": "int8_per_channel_static"} + else: + quantization_config = { + "linear_scheme_override": "int8_per_channel_weight_only" + } + llm_kwargs.update( + quantization="online", + quantization_config=quantization_config, + ) + elif args.worker_setting == "bnb-nf4" and args.worker_inflight_nf4: + llm_kwargs["quantization"] = "bitsandbytes" + elif args.worker_setting.startswith("torchao-"): + from torchao.core.config import config_to_dict + from torchao.quantization import ( + Int4WeightOnlyConfig, + Int8WeightOnlyConfig, + ) + + if args.worker_setting == "torchao-int8": + torchao_config = Int8WeightOnlyConfig() + else: + if args.dtype not in ("bfloat16", "bf16"): + raise ValueError("torchao-int4 requires --dtype bfloat16") + torchao_config = Int4WeightOnlyConfig( + group_size=128, + int4_packing_format="tile_packed_to_4d", + ) + llm_kwargs.update( + quantization="torchao", + hf_overrides={ + "quantization_config_dict_json": json.dumps( + config_to_dict(torchao_config) + ) + }, + ) + + load_started = time.perf_counter() + llm = LLM( + model=args.worker_model, + tokenizer=args.tokenizer or args.model, + dtype=args.dtype, + trust_remote_code=True, + enforce_eager=args.enforce_eager, + async_scheduling=args.async_scheduling, + enable_chunked_prefill=True, + max_num_batched_tokens=args.max_num_batched_tokens, + max_model_len=args.max_model_len, + gpu_memory_utilization=args.gpu_memory_utilization, + **llm_kwargs, + ) + load_s = time.perf_counter() - load_started + prompts = load_prompts(args) + sampling_params = SamplingParams( + temperature=0, + max_tokens=args.max_tokens, + logprobs=args.logprobs, + ignore_eos=args.ignore_eos, + ) + for _ in range(args.warmup_runs): + llm.generate(prompts, sampling_params, use_tqdm=False) + + samples = [] + signatures = [] + outputs = None + for _ in range(args.repeats): + started = time.perf_counter() + outputs = llm.generate(prompts, sampling_params, use_tqdm=False) + samples.append(time.perf_counter() - started) + signatures.append([tuple(item.outputs[0].token_ids) for item in outputs]) + assert outputs is not None + repeat_mismatch = None + for run_index, signature in enumerate(signatures[1:], start=1): + if signature == signatures[0]: + continue + for request_index, (reference_ids, candidate_ids) in enumerate( + zip(signatures[0], signature) + ): + if reference_ids == candidate_ids: + continue + common = min(len(reference_ids), len(candidate_ids)) + token_index = next( + ( + index + for index in range(common) + if reference_ids[index] != candidate_ids[index] + ), + common, + ) + repeat_mismatch = { + "run": run_index, + "request": request_index, + "token": token_index, + } + break + if repeat_mismatch is None: + repeat_mismatch = {"run": run_index, "request": None, "token": None} + break + if repeat_mismatch is not None and args.require_repeatable: + raise RuntimeError( + f"{args.worker_setting} output changed across repeated runs: " + f"{repeat_mismatch}" + ) + + elapsed = statistics.median(samples) + output_tokens = sum(len(item.outputs[0].token_ids) for item in outputs) + result = { + "setting": args.worker_setting, + "model": args.worker_model, + "load_s": load_s, + "elapsed_s": elapsed, + "samples_s": samples, + "output_tokens": output_tokens, + "output_tok_s": output_tokens / elapsed, + "repeatable": repeat_mismatch is None, + "repeat_mismatch": repeat_mismatch, + "requests": [list(item.outputs[0].token_ids) for item in outputs], + } + print("RESULT_JSON " + json.dumps(result), flush=True) + + +def _setting_model(args: argparse.Namespace, setting: str) -> tuple[str, bool]: + if setting == "fp16": + return args.model, False + if setting == "online-int8" or setting.startswith("torchao-"): + return args.model, False + if setting == "bnb-int8": + if args.int8_model is None: + raise ValueError("--int8-model is required for bnb-int8") + return args.int8_model, False + if args.int4_model is None: + return args.model, True + return args.int4_model, False + + +def run_setting(args: argparse.Namespace, setting: str) -> dict[str, Any]: + model, inflight_nf4 = _setting_model(args, setting) + command = [ + sys.executable, + __file__, + "--model", + args.model, + "--worker-model", + model, + "--worker-setting", + setting, + ] + for name in ( + "max_tokens", + "logprobs", + "warmup_runs", + "repeats", + "dtype", + "max_model_len", + "max_num_batched_tokens", + "gpu_memory_utilization", + ): + command.extend(["--" + name.replace("_", "-"), str(getattr(args, name))]) + if args.tokenizer is not None: + command.extend(["--tokenizer", args.tokenizer]) + for prompt in args.prompt: + command.extend(["--prompt", prompt]) + if args.prompts_file is not None: + command.extend(["--prompts-file", str(args.prompts_file)]) + command.append("--enforce-eager" if args.enforce_eager else "--no-enforce-eager") + command.append( + "--async-scheduling" if args.async_scheduling else "--no-async-scheduling" + ) + command.append("--ignore-eos" if args.ignore_eos else "--no-ignore-eos") + command.append( + "--require-repeatable" if args.require_repeatable else "--no-require-repeatable" + ) + if inflight_nf4: + command.append("--worker-inflight-nf4") + + result = None + model_memory_gib = None + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + assert process.stdout is not None + for line in process.stdout: + if match := MODEL_MEMORY_PATTERN.search(line): + model_memory_gib = float(match.group(1)) + if line.startswith("RESULT_JSON "): + result = json.loads(line.removeprefix("RESULT_JSON ")) + elif not any(message in line for message in NOISY_BNB_MESSAGES): + print(f"[{setting}] {line}", end="") + return_code = process.wait() + if return_code != 0: + raise RuntimeError(f"{setting} worker exited with status {return_code}") + if result is None: + raise RuntimeError(f"{setting} worker produced no result") + result["model_memory_gib"] = model_memory_gib + return result + + +def compare_with_fp16( + reference: dict[str, Any], candidate: dict[str, Any] +) -> dict[str, Any]: + if len(reference["requests"]) != len(candidate["requests"]): + raise RuntimeError("FP16 and quantized settings returned different batch sizes") + + matching_tokens = 0 + total_tokens = 0 + exact_requests = 0 + request_reports = [] + for index, (ref_ids, candidate_ids) in enumerate( + zip(reference["requests"], candidate["requests"]) + ): + common = min(len(ref_ids), len(candidate_ids)) + first_difference = next( + (pos for pos in range(common) if ref_ids[pos] != candidate_ids[pos]), + None, + ) + if first_difference is None and len(ref_ids) != len(candidate_ids): + first_difference = common + exact = first_difference is None + exact_requests += int(exact) + matching_tokens += sum( + ref_token == candidate_token + for ref_token, candidate_token in zip(ref_ids, candidate_ids) + ) + total_tokens += max(len(ref_ids), len(candidate_ids)) + request_reports.append( + { + "request": index, + "exact": exact, + "first_difference": first_difference, + "fp16_length": len(ref_ids), + "candidate_length": len(candidate_ids), + } + ) + + memory_reduction = None + if ( + reference["model_memory_gib"] is not None + and candidate["model_memory_gib"] is not None + ): + memory_reduction = 1.0 - ( + candidate["model_memory_gib"] / reference["model_memory_gib"] + ) + return { + "setting": candidate["setting"], + "speed_ratio": candidate["output_tok_s"] / reference["output_tok_s"], + "memory_reduction": memory_reduction, + "token_agreement": matching_tokens / total_tokens if total_tokens else 1.0, + "exact_requests": exact_requests, + "total_requests": len(request_reports), + "fp16_output_tok_s": reference["output_tok_s"], + "candidate_output_tok_s": candidate["output_tok_s"], + "fp16_model_memory_gib": reference["model_memory_gib"], + "candidate_model_memory_gib": candidate["model_memory_gib"], + "requests": request_reports, + } + + +def main() -> None: + args = parse_args() + if args.worker_setting is not None: + run_worker(args) + return + if "fp16" not in args.settings: + raise ValueError("--settings must include fp16 as the reference") + + results = {setting: run_setting(args, setting) for setting in args.settings} + comparisons = [ + compare_with_fp16(results["fp16"], results[setting]) + for setting in args.settings + if setting != "fp16" + ] + report = {"results": results, "comparisons": comparisons} + print(json.dumps(report, indent=2)) + + failed = [] + for comparison in comparisons: + if comparison["speed_ratio"] < args.min_speed_ratio: + failed.append( + f"{comparison['setting']} speed_ratio={comparison['speed_ratio']:.4f}" + ) + reduction = comparison["memory_reduction"] + if reduction is None or reduction < args.min_memory_reduction: + failed.append(f"{comparison['setting']} memory_reduction={reduction}") + if args.require_gates and failed: + print("FAILED_GATES " + "; ".join(failed), file=sys.stderr) + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_rwkv7.py b/benchmarks/kernels/benchmark_rwkv7.py new file mode 100644 index 000000000000..d5b0200eeb2e --- /dev/null +++ b/benchmarks/kernels/benchmark_rwkv7.py @@ -0,0 +1,291 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare an RWKV7 candidate kernel policy with Torch through the vLLM engine. + +Each backend runs in a fresh subprocess so environment selection, CUDA graphs, +and recurrent caches cannot leak between runs. The report identifies the first +greedy-token divergence and includes the top log-probabilities at that position. +""" + +import argparse +import json +import os +import statistics +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +DEFAULT_PROMPTS = [ + "The Eiffel Tower is located in", + "A short proof that there are infinitely many primes begins", + "Write a Python function that computes Fibonacci numbers:", + "The most important property of recurrent neural networks is", + "Once upon a time in a quiet village", + "Explain why the sky appears blue during the day.", + "In numerical analysis, floating point accumulation order", + "Translate to Chinese: artificial intelligence inference engine", +] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True) + parser.add_argument("--tokenizer") + parser.add_argument("--prompt", action="append", default=[]) + parser.add_argument("--prompts-file", type=Path) + parser.add_argument("--max-tokens", type=int, default=64) + parser.add_argument("--logprobs", type=int, default=5) + parser.add_argument("--warmup-runs", type=int, default=1) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--dtype", default="auto") + parser.add_argument("--max-model-len", type=int, default=2048) + parser.add_argument("--max-num-batched-tokens", type=int, default=32) + parser.add_argument("--gpu-memory-utilization", type=float, default=0.8) + parser.add_argument( + "--enforce-eager", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--async-scheduling", action=argparse.BooleanOptionalAction, default=False + ) + parser.add_argument( + "--ignore-eos", action=argparse.BooleanOptionalAction, default=False + ) + parser.add_argument( + "--require-exact", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--candidate-backend", choices=("auto", "triton"), default="triton" + ) + parser.add_argument( + "--worker", + choices=("auto", "torch", "triton"), + help=argparse.SUPPRESS, + ) + return parser.parse_args() + + +def load_prompts(args: argparse.Namespace) -> list[str]: + if args.warmup_runs < 0 or args.repeats < 1: + raise ValueError("--warmup-runs must be >= 0 and --repeats must be >= 1") + prompts = list(args.prompt) + if args.prompts_file is not None: + loaded = json.loads(args.prompts_file.read_text()) + if not isinstance(loaded, list) or not all(isinstance(x, str) for x in loaded): + raise ValueError("--prompts-file must contain a JSON array of strings") + prompts.extend(loaded) + return prompts or DEFAULT_PROMPTS + + +def serialize_logprobs(logprobs: Any) -> list[list[dict[str, Any]]]: + serialized = [] + for position in logprobs or []: + candidates = [ + { + "token_id": int(token_id), + "logprob": float(value.logprob), + "rank": value.rank, + } + for token_id, value in position.items() + ] + candidates.sort(key=lambda item: item["logprob"], reverse=True) + serialized.append(candidates) + return serialized + + +def run_worker(args: argparse.Namespace) -> None: + assert args.worker is not None + os.environ["VLLM_RWKV7_KERNEL"] = args.worker + + from vllm import LLM, SamplingParams + + prompts = load_prompts(args) + llm = LLM( + model=args.model, + tokenizer=args.tokenizer or args.model, + dtype=args.dtype, + trust_remote_code=True, + enforce_eager=args.enforce_eager, + async_scheduling=args.async_scheduling, + enable_chunked_prefill=True, + max_num_batched_tokens=args.max_num_batched_tokens, + max_model_len=args.max_model_len, + gpu_memory_utilization=args.gpu_memory_utilization, + ) + sampling_params = SamplingParams( + temperature=0, + max_tokens=args.max_tokens, + logprobs=args.logprobs, + ignore_eos=args.ignore_eos, + ) + for _ in range(args.warmup_runs): + llm.generate(prompts, sampling_params, use_tqdm=False) + samples = [] + signatures = [] + outputs = None + for _ in range(args.repeats): + started = time.perf_counter() + outputs = llm.generate(prompts, sampling_params, use_tqdm=False) + samples.append(time.perf_counter() - started) + signatures.append([tuple(item.outputs[0].token_ids) for item in outputs]) + assert outputs is not None + if any(signature != signatures[0] for signature in signatures[1:]): + raise RuntimeError(f"{args.worker} output changed across repeated runs") + elapsed = statistics.median(samples) + result = { + "backend": args.worker, + "elapsed_s": elapsed, + "samples_s": samples, + "output_tokens": sum(len(item.outputs[0].token_ids) for item in outputs), + "requests": [ + { + "token_ids": list(item.outputs[0].token_ids), + "logprobs": serialize_logprobs(item.outputs[0].logprobs), + } + for item in outputs + ], + } + result["output_tok_s"] = result["output_tokens"] / elapsed + print("RESULT_JSON " + json.dumps(result), flush=True) + + +def run_backend(args: argparse.Namespace, backend: str) -> dict[str, Any]: + command = [sys.executable, __file__] + for name in ( + "model", + "max_tokens", + "logprobs", + "warmup_runs", + "repeats", + "dtype", + "max_model_len", + "max_num_batched_tokens", + "gpu_memory_utilization", + ): + command.extend(["--" + name.replace("_", "-"), str(getattr(args, name))]) + if args.tokenizer is not None: + command.extend(["--tokenizer", args.tokenizer]) + for prompt in args.prompt: + command.extend(["--prompt", prompt]) + if args.prompts_file is not None: + command.extend(["--prompts-file", str(args.prompts_file)]) + command.append("--enforce-eager" if args.enforce_eager else "--no-enforce-eager") + command.append( + "--async-scheduling" if args.async_scheduling else "--no-async-scheduling" + ) + command.append("--ignore-eos" if args.ignore_eos else "--no-ignore-eos") + command.extend(["--worker", backend]) + + result = None + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + assert process.stdout is not None + for line in process.stdout: + if line.startswith("RESULT_JSON "): + result = json.loads(line.removeprefix("RESULT_JSON ")) + else: + print(f"[{backend}] {line}", end="") + return_code = process.wait() + if return_code != 0: + raise RuntimeError(f"{backend} worker exited with status {return_code}") + if result is None: + raise RuntimeError(f"{backend} worker produced no result") + return result + + +def top_margin(candidates: list[dict[str, Any]]) -> float | None: + if len(candidates) < 2: + return None + return candidates[0]["logprob"] - candidates[1]["logprob"] + + +def compare_results( + torch_result: dict[str, Any], candidate_result: dict[str, Any] +) -> dict[str, Any]: + if len(torch_result["requests"]) != len(candidate_result["requests"]): + raise RuntimeError("Reference and candidate returned different batch sizes") + request_reports = [] + exact_requests = 0 + for request_idx, (torch_request, candidate_request) in enumerate( + zip(torch_result["requests"], candidate_result["requests"]) + ): + torch_ids = torch_request["token_ids"] + candidate_ids = candidate_request["token_ids"] + first_difference = next( + ( + index + for index, (torch_id, candidate_id) in enumerate( + zip(torch_ids, candidate_ids) + ) + if torch_id != candidate_id + ), + None, + ) + if first_difference is None and len(torch_ids) != len(candidate_ids): + first_difference = min(len(torch_ids), len(candidate_ids)) + exact = first_difference is None + exact_requests += int(exact) + report: dict[str, Any] = { + "request": request_idx, + "exact": exact, + "torch_length": len(torch_ids), + "candidate_length": len(candidate_ids), + "first_difference": first_difference, + } + if first_difference is not None: + position = first_difference + torch_top = ( + torch_request["logprobs"][position] + if position < len(torch_request["logprobs"]) + else [] + ) + candidate_top = ( + candidate_request["logprobs"][position] + if position < len(candidate_request["logprobs"]) + else [] + ) + report.update( + torch_token=torch_ids[position] if position < len(torch_ids) else None, + candidate_token=( + candidate_ids[position] if position < len(candidate_ids) else None + ), + torch_top_logprobs=torch_top, + candidate_top_logprobs=candidate_top, + torch_top2_margin=top_margin(torch_top), + candidate_top2_margin=top_margin(candidate_top), + ) + request_reports.append(report) + return { + "exact": exact_requests == len(request_reports), + "exact_requests": exact_requests, + "total_requests": len(request_reports), + "candidate_backend": candidate_result["backend"], + "torch_output_tok_s": torch_result["output_tok_s"], + "candidate_output_tok_s": candidate_result["output_tok_s"], + "speedup": candidate_result["output_tok_s"] / torch_result["output_tok_s"], + "requests": request_reports, + } + + +def main() -> None: + args = parse_args() + if args.worker is not None: + run_worker(args) + return + + torch_result = run_backend(args, "torch") + candidate_result = run_backend(args, args.candidate_backend) + comparison = compare_results(torch_result, candidate_result) + print(json.dumps(comparison, indent=2)) + if args.require_exact and not comparison["exact"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md index a4da52557afb..b758a97a3477 100644 --- a/docs/features/quantization/online.md +++ b/docs/features/quantization/online.md @@ -37,6 +37,36 @@ vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8 | `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance | | `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | | | `mxfp8` | fp8_e4m3 data, e8m0 per-1x32-block scale | fp8_e4m3 data, e8m0 per-1x32-block scale | Requires SM 100+ (Blackwell or newer) for w8a8, other GPUs use a w8a16 fallback | +| `int8_per_channel_static` | int8 data, fp32 per-output-channel scale | int8 data, fp32 dynamic per-token scale | Configure explicitly through `quantization_config.linear`; requires compute capability 7.5 or newer | + +### Dense INT8 + +Dense linear layers can be quantized from an FP16/BF16 checkpoint at load time. +The checkpoint itself does not need to be converted: + +```python +from vllm import LLM + +llm = LLM( + "RWKV/RWKV7-Goose-World2.8-1.5B-HF", + quantization="online", + quantization_config={"linear": "int8_per_channel_static"}, +) +``` + +The same configuration can be supplied to `vllm serve`: + +```bash +vllm serve RWKV/RWKV7-Goose-World2.8-1.5B-HF \ + --quantization online \ + --quantization-config.linear int8_per_channel_static +``` + +RWKV7 keeps its recurrent key/value and low-rank control projections in the +requested model dtype. The larger receptance, output, feed-forward, and untied +LM-head projections use the configured quantization method. This avoids +compounding quantization error in the recurrent state while retaining most of +the weight-memory reduction. ## Advanced Configuration diff --git a/docs/features/quantization/torchao.md b/docs/features/quantization/torchao.md index b95b560882bb..e4291944b6c1 100644 --- a/docs/features/quantization/torchao.md +++ b/docs/features/quantization/torchao.md @@ -41,3 +41,33 @@ You can quantize your own huggingface model with torchao, e.g. [transformers](ht ``` Alternatively, you can use the [TorchAO Quantization space](https://huggingface.co/spaces/medmekk/TorchAO_Quantization) for quantizing models with a simple UI. + +## Quantizing at Load Time + +TorchAO can also quantize an FP16/BF16 checkpoint while vLLM loads it. For +example, the following uses the packed INT4 weight-only kernel: + +```python +import json + +from torchao.core.config import config_to_dict +from torchao.quantization import Int4WeightOnlyConfig +from vllm import LLM + +config = Int4WeightOnlyConfig( + group_size=128, + int4_packing_format="tile_packed_to_4d", +) +llm = LLM( + "RWKV/RWKV7-Goose-World2.8-1.5B-HF", + dtype="bfloat16", + quantization="torchao", + hf_overrides={ + "quantization_config_dict_json": json.dumps(config_to_dict(config)) + }, +) +``` + +For models with an untied output embedding, TorchAO also quantizes the +`ParallelLMHead`; leaving a large LM head unquantized can otherwise dominate +both memory use and decode time. diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0836f11e267d..263fb5a42b83 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -409,6 +409,7 @@ th { | `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ | | `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ | | `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ | +| `RWKV7ForCausalLM` | RWKV7 | `RWKV/RWKV7-Goose-World2.8-0.1B-HF`, etc. | | ✅︎ | | `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ | | `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ | | `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ | diff --git a/tests/model_executor/test_rwkv7.py b/tests/model_executor/test_rwkv7.py new file mode 100644 index 000000000000..349bb122894a --- /dev/null +++ b/tests/model_executor/test_rwkv7.py @@ -0,0 +1,1270 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import sys +import tempfile +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import torch +from transformers import AutoTokenizer + +from vllm.config import ( + CacheConfig, + CompilationConfig, + CompilationMode, + DeviceConfig, + ModelConfig, + ParallelConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.distributed.parallel_state import ( + ensure_model_parallel_initialized, + init_distributed_environment, +) +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.linear import UnquantizedLinearMethod +from vllm.model_executor.layers.mamba.mamba_utils import ( + get_conv_copy_spec, + get_temporal_copy_spec, +) +from vllm.model_executor.layers.rwkv7 import ( + RWKV7KernelBackend, + diagnose_rwkv7_recurrent_scan_packed, + resolve_rwkv7_kernel_backend, + rwkv7_recurrent_scan_packed, +) +from vllm.model_executor.models.config import MODELS_CONFIG_MAP, MambaModelConfig +from vllm.model_executor.models.rwkv7 import ( + RWKV7Attention, + RWKV7Block, + RWKV7ForCausalLM, +) +from vllm.transformers_utils.configs.rwkv7 import RWKV7Config +from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata +from vllm.v1.attention.backends.utils import PAD_SLOT_ID + +try: + import pytest +except ImportError: + pytest = None + + +def _make_config() -> RWKV7Config: + return RWKV7Config( + vocab_size=128, + hidden_size=64, + hidden_ratio=2, + num_hidden_layers=2, + head_dim=16, + num_heads=4, + decay_low_rank_dim=16, + gate_low_rank_dim=16, + a_low_rank_dim=16, + v_low_rank_dim=16, + norm_bias=True, + value_dim=64, + ) + + +def _initialize_module_parameters(module: torch.nn.Module) -> None: + generator = torch.Generator().manual_seed(0) + for name, parameter in module.named_parameters(): + if parameter.ndim == 0 or name.endswith(".bias"): + parameter.data.zero_() + elif "g_norm.weight" in name or "k_a" in name: + parameter.data.fill_(1.0) + else: + parameter.data.normal_(mean=0.0, std=0.02, generator=generator) + + +def _rwkv7_recurrent_oracle_packed( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + state: torch.Tensor, + query_start_loc: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + output = torch.empty_like(v) + final_state = state.clone() + for seq_idx in range(query_start_loc.numel() - 1): + start = int(query_start_loc[seq_idx].item()) + end = int(query_start_loc[seq_idx + 1].item()) + seq_state = state[seq_idx].transpose(-2, -1) + for token_idx in range(start, end): + sa = (seq_state * (-kk[token_idx]).unsqueeze(-1)).sum(dim=-2) + seq_state = ( + torch.exp(w[token_idx]).unsqueeze(-1) * seq_state + + (kk[token_idx] * a[token_idx]).unsqueeze(-1) * sa.unsqueeze(-2) + + k[token_idx].unsqueeze(-1) * v[token_idx].unsqueeze(-2) + ) + output[token_idx] = (seq_state * r[token_idx].unsqueeze(-1)).sum(dim=-2) + final_state[seq_idx] = seq_state.transpose(-2, -1) + return output, final_state + + +def _test_rwkv7_recurrent_scan_packed( + device: str, *, backend: RWKV7KernelBackend +) -> None: + generator = torch.Generator(device=device).manual_seed(123) + num_tokens, num_heads, key_dim, value_dim = 7, 3, 5, 4 + vectors = [ + torch.randn( + num_tokens, + num_heads, + key_dim, + device=device, + generator=generator, + ) + for _ in range(5) + ] + r, w, k, kk, a = vectors + w = -w.abs() + v = torch.randn( + num_tokens, + num_heads, + value_dim, + device=device, + generator=generator, + ) + state = torch.randn( + 2, + num_heads, + value_dim, + key_dim, + device=device, + generator=generator, + ) + query_start_loc = torch.tensor([0, 2, 7], dtype=torch.int32, device=device) + + output, final_state = rwkv7_recurrent_scan_packed( + r, w, k, v, kk, a, state, query_start_loc, backend=backend + ) + expected_output, expected_state = _rwkv7_recurrent_oracle_packed( + r, w, k, v, kk, a, state, query_start_loc + ) + + auto_output, auto_state = rwkv7_recurrent_scan_packed( + r, w, k, v, kk, a, state, query_start_loc, backend="auto" + ) + torch_output, torch_state = rwkv7_recurrent_scan_packed( + r, w, k, v, kk, a, state, query_start_loc, backend="torch" + ) + torch.testing.assert_close(auto_output, torch_output, rtol=0, atol=0) + torch.testing.assert_close(auto_state, torch_state, rtol=0, atol=0) + + torch.testing.assert_close(output, expected_output, rtol=2e-4, atol=1e-4) + torch.testing.assert_close(final_state, expected_state, rtol=2e-4, atol=1e-4) + + if backend == "triton": + report = diagnose_rwkv7_recurrent_scan_packed( + r, + w, + k, + v, + kk, + a, + state, + query_start_loc, + rtol=2e-4, + atol=1e-4, + ) + assert report.output_close + assert report.state_close + assert report.first_output_mismatch_token is None + assert report.first_state_mismatch_sequence is None + + +def test_rwkv7_recurrent_scan_packed_matches_oracle(): + _test_rwkv7_recurrent_scan_packed("cpu", backend="torch") + + +def test_rwkv7_recurrent_scan_packed_matches_oracle_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required to exercise the RWKV7 Triton scan.") + _test_rwkv7_recurrent_scan_packed("cuda", backend="triton") + + +def test_rwkv7_kernel_auto_is_fail_closed_on_cpu(monkeypatch): + monkeypatch.delenv("VLLM_RWKV7_KERNEL", raising=False) + tensor = torch.empty(1) + assert resolve_rwkv7_kernel_backend() == "torch" + assert ( + resolve_rwkv7_kernel_backend("auto", input_tensor=tensor, state=tensor.float()) + == "torch" + ) + assert resolve_rwkv7_kernel_backend("torch") == "torch" + + +def test_rwkv7_explicit_triton_rejects_cpu(): + tensor = torch.empty(1) + with pytest.raises(RuntimeError, match="RWKV7 Triton kernel"): + resolve_rwkv7_kernel_backend( + "triton", input_tensor=tensor, state=tensor.float() + ) + + +def test_rwkv7_kernel_environment_rejects_invalid_value(monkeypatch): + monkeypatch.setenv("VLLM_RWKV7_KERNEL", "invalid") + with pytest.raises(ValueError, match="VLLM_RWKV7_KERNEL"): + resolve_rwkv7_kernel_backend() + + +def _make_prefill_metadata( + seq_len: int, *, device: torch.device +) -> LinearAttentionMetadata: + return LinearAttentionMetadata( + num_prefills=1, + num_prefill_tokens=seq_len, + num_decodes=0, + num_decode_tokens=0, + query_start_loc=torch.tensor([0, seq_len], dtype=torch.int32, device=device), + seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=device), + state_indices_tensor=torch.tensor([0], dtype=torch.long, device=device), + ) + + +def _make_decode_metadata( + total_seq_len: int, *, device: torch.device +) -> LinearAttentionMetadata: + return LinearAttentionMetadata( + num_prefills=0, + num_prefill_tokens=0, + num_decodes=1, + num_decode_tokens=1, + query_start_loc=torch.tensor([0, 1], dtype=torch.int32, device=device), + seq_lens=torch.tensor([total_seq_len], dtype=torch.int32, device=device), + state_indices_tensor=torch.tensor([0], dtype=torch.long, device=device), + ) + + +def _make_multi_decode_metadata( + total_seq_lens: list[int], state_indices: list[int], *, device: torch.device +) -> LinearAttentionMetadata: + num_decodes = len(total_seq_lens) + return LinearAttentionMetadata( + num_prefills=0, + num_prefill_tokens=0, + num_decodes=num_decodes, + num_decode_tokens=num_decodes, + query_start_loc=torch.arange( + 0, num_decodes + 1, dtype=torch.int32, device=device + ), + seq_lens=torch.tensor(total_seq_lens, dtype=torch.int32, device=device), + state_indices_tensor=torch.tensor( + state_indices, dtype=torch.long, device=device + ), + ) + + +def _make_multi_prefill_metadata( + query_lens: list[int], + total_seq_lens: list[int], + state_indices: list[int], + *, + device: torch.device, +) -> LinearAttentionMetadata: + query_start_loc = [0] + for query_len in query_lens: + query_start_loc.append(query_start_loc[-1] + query_len) + return LinearAttentionMetadata( + num_prefills=len(query_lens), + num_prefill_tokens=query_start_loc[-1], + num_decodes=0, + num_decode_tokens=0, + query_start_loc=torch.tensor( + query_start_loc, + dtype=torch.int32, + device=device, + ), + seq_lens=torch.tensor(total_seq_lens, dtype=torch.int32, device=device), + state_indices_tensor=torch.tensor( + state_indices, dtype=torch.long, device=device + ), + ) + + +def _require_reference_checkpoint() -> tuple[Path, Any]: + """Resolve optional dependencies for RWKV7 reference parity tests only. + + The external top-level ``fla`` package is only needed here so parity tests + can import the reference Hugging Face implementation from a + flash-linear-attention checkout. RWKV7 runtime in this PR stays within the + vLLM tree and does not depend on that external package. + """ + if pytest is None: + raise RuntimeError("pytest is required to run RWKV7 integration tests.") + + model_path = os.getenv("VLLM_RWKV7_TEST_MODEL_PATH") + fla_path = os.getenv("VLLM_RWKV7_TEST_FLA_PATH") + + if not model_path: + pytest.skip( + "Set VLLM_RWKV7_TEST_MODEL_PATH to run optional RWKV7 reference " + "parity tests." + ) + if not fla_path: + pytest.skip( + "Set VLLM_RWKV7_TEST_FLA_PATH to a flash-linear-attention checkout " + "to run optional RWKV7 reference parity tests. vLLM runtime does " + "not depend on the external top-level `fla` package." + ) + + assert model_path + assert fla_path + model_dir = Path(model_path) + fla_dir = Path(fla_path) + if not model_dir.exists(): + pytest.skip(f"RWKV7 model path does not exist: {model_dir}") + if not fla_dir.exists(): + pytest.skip(f"FLA path does not exist: {fla_dir}") + + if str(fla_dir) not in sys.path: + sys.path.insert(0, str(fla_dir)) + + # Test-only import of the reference implementation from an external FLA + # checkout. RWKV7 runtime itself does not depend on that package. + from fla.models.rwkv7 import RWKV7ForCausalLM as ReferenceRWKV7ForCausalLM + + return model_dir, ReferenceRWKV7ForCausalLM + + +def _write_test_model_config(tmp_path: Path, config: RWKV7Config | None = None) -> Path: + model_dir = tmp_path / "rwkv7-test-model" + config = _make_config() if config is None else config + config.architectures = ["RWKV7ForCausalLM"] + config.save_pretrained(model_dir) + return model_dir + + +def _make_vllm_config( + model_path: Path, + *, + dtype: str = "float32", + device: str = "cuda", +) -> VllmConfig: + return VllmConfig( + model_config=ModelConfig( + str(model_path), + trust_remote_code=False, + dtype=dtype, + runner="generate", + ), + parallel_config=ParallelConfig( + tensor_parallel_size=1, + pipeline_parallel_size=1, + ), + cache_config=CacheConfig(), + compilation_config=CompilationConfig(mode=CompilationMode.NONE), + device_config=DeviceConfig(device), + ) + + +def _local_distributed_init_method() -> str: + os.environ["MASTER_ADDR"] = "localhost" + os.environ["GLOO_SOCKET_IFNAME"] = "lo" + os.environ.setdefault("RANK", "0") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + fd, path = tempfile.mkstemp(prefix="rwkv7-dist-") + os.close(fd) + return f"file://{path}" + + +def _allocate_kv_cache(model: RWKV7ForCausalLM, *, device: torch.device) -> None: + for layer in model.model.layers: + state_shapes = layer.get_state_shape() + state_dtypes = layer.get_state_dtype() + layer.kv_cache = tuple( + torch.zeros((1, *shape), dtype=dtype, device=device) + for shape, dtype in zip(state_shapes, state_dtypes) + ) + + +def test_rwkv7_block_forward_without_metadata(): + config = _make_config() + vllm_config = VllmConfig(device_config=DeviceConfig("cpu")) + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1, backend="gloo") + try: + block0 = RWKV7Block(config=config, layer_idx=0, prefix="model.layers.0") + block1 = RWKV7Block(config=config, layer_idx=1, prefix="model.layers.1") + _initialize_module_parameters(block0) + _initialize_module_parameters(block1) + + hidden_states = torch.randn(5, config.hidden_size) + hidden_states, v_first = block0(hidden_states, None, None) + hidden_states, v_first = block1(hidden_states, v_first, None) + + assert hidden_states.shape == (5, config.hidden_size) + assert v_first.shape == (5, config.hidden_size) + assert torch.isfinite(hidden_states).all() + assert torch.isfinite(v_first).all() + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_block_registers_static_forward_context(): + config = _make_config() + vllm_config = VllmConfig(device_config=DeviceConfig("cpu")) + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1, backend="gloo") + try: + prefix = "model.layers.0" + block = RWKV7Block(config=config, layer_idx=0, prefix=prefix) + assert ( + vllm_config.compilation_config.static_forward_context[prefix] is block + ) + assert ( + vllm_config.compilation_config.static_forward_context[f"{prefix}.attn"] + is block.attn + ) + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_attention_custom_op_matches_direct_forward(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required to exercise the RWKV7 attention custom op.") + + config = _make_config() + vllm_config = VllmConfig(device_config=DeviceConfig("cuda")) + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="nccl", + ) + ensure_model_parallel_initialized(1, 1, backend="nccl") + try: + block = RWKV7Block(config=config, layer_idx=0, prefix="model.layers.0") + _initialize_module_parameters(block) + block = block.to("cuda", torch.float32) + hidden_states = torch.randn(4, config.hidden_size, device="cuda") + + direct = block.attn._forward(hidden_states, None, None, None) + with set_forward_context(None, vllm_config): + wrapped = block.attn(hidden_states, None, None, None) + + for wrapped_tensor, direct_tensor in zip(wrapped, direct): + torch.testing.assert_close(wrapped_tensor, direct_tensor) + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_block_updates_cached_states(): + config = _make_config() + vllm_config = VllmConfig(device_config=DeviceConfig("cpu")) + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1, backend="gloo") + try: + block = RWKV7Block(config=config, layer_idx=0, prefix="model.layers.0") + _initialize_module_parameters(block) + + block.kv_cache = ( + torch.zeros(1, config.hidden_size), + torch.zeros(1, config.num_heads, config.head_dim, config.head_dim), + torch.zeros(1, config.hidden_size), + ) + + prefill_metadata = _make_prefill_metadata(3, device=torch.device("cpu")) + hidden_states = torch.randn(3, config.hidden_size) + output, v_first = block(hidden_states, None, prefill_metadata) + + assert output.shape == hidden_states.shape + assert v_first.shape == hidden_states.shape + assert torch.isfinite(output).all() + assert block.kv_cache[0][0].abs().sum() > 0 + assert block.kv_cache[1][0].abs().sum() > 0 + assert block.kv_cache[2][0].abs().sum() > 0 + + decode_metadata = _make_decode_metadata(4, device=torch.device("cpu")) + decode_hidden = torch.randn(1, config.hidden_size) + decode_output, decode_v_first = block( + decode_hidden, v_first[:1].clone(), decode_metadata + ) + + assert decode_output.shape == decode_hidden.shape + assert decode_v_first.shape == decode_hidden.shape + assert torch.isfinite(decode_output).all() + assert torch.isfinite(decode_v_first).all() + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_block_batches_decode_tokens_without_changing_results(): + config = _make_config() + vllm_config = VllmConfig(device_config=DeviceConfig("cpu")) + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1, backend="gloo") + try: + block_batched = RWKV7Block( + config=config, layer_idx=0, prefix="model.layers.0" + ) + _initialize_module_parameters(block_batched) + block_ref = RWKV7Block(config=config, layer_idx=0, prefix="model.layers.1") + block_ref.load_state_dict(block_batched.state_dict()) + + generator = torch.Generator().manual_seed(123) + state_shapes = block_batched.get_state_shape() + state_dtypes = block_batched.get_state_dtype() + + def make_cache() -> tuple[torch.Tensor, ...]: + return tuple( + torch.randn( + (2, *shape), + generator=generator, + dtype=dtype, + ) + for shape, dtype in zip(state_shapes, state_dtypes) + ) + + block_batched.kv_cache = make_cache() + block_ref.kv_cache = tuple( + cache.clone() for cache in block_batched.kv_cache + ) + + hidden_states = torch.randn( + 2, config.hidden_size, generator=generator, dtype=torch.float32 + ) + metadata = _make_multi_decode_metadata( + [5, 7], [0, 1], device=torch.device("cpu") + ) + + output_batched, v_first_batched = block_batched( + hidden_states, None, metadata + ) + + output_ref = torch.empty_like(hidden_states) + v_first_ref = torch.empty_like(hidden_states) + for idx, slot_id in enumerate([0, 1]): + states = block_ref._get_kv_state(slot_id, use_initial_state=True) + ( + out, + v_first_out, + attn_shift, + recurrent, + ffn_shift, + ) = block_ref._run_sequence( + hidden_states[idx : idx + 1], + None, + *states, + ) + output_ref[idx : idx + 1] = out + v_first_ref[idx : idx + 1] = v_first_out + block_ref._store_kv_state(slot_id, attn_shift, recurrent, ffn_shift) + + torch.testing.assert_close(output_batched, output_ref) + torch.testing.assert_close(v_first_batched, v_first_ref) + for batched_state, ref_state in zip( + block_batched.kv_cache, block_ref.kv_cache + ): + torch.testing.assert_close(batched_state, ref_state) + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_block_ignores_padded_cudagraph_decode_slots(): + config = _make_config() + vllm_config = VllmConfig(device_config=DeviceConfig("cpu")) + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1, backend="gloo") + try: + block_padded = RWKV7Block( + config=config, layer_idx=0, prefix="model.layers.0" + ) + _initialize_module_parameters(block_padded) + block_ref = RWKV7Block(config=config, layer_idx=0, prefix="model.layers.1") + block_ref.load_state_dict(block_padded.state_dict()) + + generator = torch.Generator().manual_seed(321) + state_shapes = block_padded.get_state_shape() + state_dtypes = block_padded.get_state_dtype() + block_padded.kv_cache = tuple( + torch.randn((2, *shape), generator=generator, dtype=dtype) + for shape, dtype in zip(state_shapes, state_dtypes) + ) + block_ref.kv_cache = tuple(state.clone() for state in block_padded.kv_cache) + hidden_states = torch.randn( + 2, config.hidden_size, generator=generator, dtype=torch.float32 + ) + + padded_metadata = LinearAttentionMetadata( + num_prefills=0, + num_prefill_tokens=0, + num_decodes=3, + num_decode_tokens=2, + query_start_loc=torch.tensor([0, 1, 2, 2], dtype=torch.int32), + seq_lens=torch.tensor([5, 7, 0], dtype=torch.int32), + state_indices_tensor=torch.tensor( + [0, 1, PAD_SLOT_ID], dtype=torch.long + ), + ) + ref_metadata = _make_multi_decode_metadata( + [5, 7], [0, 1], device=torch.device("cpu") + ) + + output_padded, v_first_padded = block_padded( + hidden_states, None, padded_metadata + ) + output_ref, v_first_ref = block_ref(hidden_states, None, ref_metadata) + + torch.testing.assert_close(output_padded, output_ref) + torch.testing.assert_close(v_first_padded, v_first_ref) + for padded_state, ref_state in zip( + block_padded.kv_cache, block_ref.kv_cache + ): + torch.testing.assert_close(padded_state, ref_state) + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_block_batches_decode_tokens_without_changing_results_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required to exercise fused RWKV7 decode batching.") + + config = _make_config() + vllm_config = VllmConfig(device_config=DeviceConfig("cuda")) + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="nccl", + ) + ensure_model_parallel_initialized(1, 1, backend="nccl") + try: + block_batched = RWKV7Block( + config=config, layer_idx=0, prefix="model.layers.0" + ) + _initialize_module_parameters(block_batched) + block_batched = block_batched.to("cuda", torch.float32) + + block_ref = RWKV7Block(config=config, layer_idx=0, prefix="model.layers.1") + block_ref.load_state_dict(block_batched.state_dict()) + block_ref = block_ref.to("cuda", torch.float32) + + torch.manual_seed(123) + state_shapes = block_batched.get_state_shape() + state_dtypes = block_batched.get_state_dtype() + + def make_cache() -> tuple[torch.Tensor, ...]: + return tuple( + torch.randn( + (2, *shape), + dtype=dtype, + device="cuda", + ) + for shape, dtype in zip(state_shapes, state_dtypes) + ) + + block_batched.kv_cache = make_cache() + block_ref.kv_cache = tuple( + cache.clone() for cache in block_batched.kv_cache + ) + + hidden_states = torch.randn( + 2, + config.hidden_size, + device="cuda", + dtype=torch.float32, + ) + metadata = _make_multi_decode_metadata( + [5, 7], [0, 1], device=torch.device("cuda") + ) + + output_batched, v_first_batched = block_batched( + hidden_states, None, metadata + ) + + output_ref = torch.empty_like(hidden_states) + v_first_ref = torch.empty_like(hidden_states) + for idx, slot_id in enumerate([0, 1]): + states = block_ref._get_kv_state(slot_id, use_initial_state=True) + ( + out, + v_first_out, + attn_shift, + recurrent, + ffn_shift, + ) = block_ref._run_sequence( + hidden_states[idx : idx + 1], + None, + *states, + ) + output_ref[idx : idx + 1] = out + v_first_ref[idx : idx + 1] = v_first_out + block_ref._store_kv_state(slot_id, attn_shift, recurrent, ffn_shift) + + torch.testing.assert_close(output_batched, output_ref, rtol=2e-4, atol=1e-3) + torch.testing.assert_close( + v_first_batched, v_first_ref, rtol=2e-4, atol=1e-3 + ) + for batched_state, ref_state in zip( + block_batched.kv_cache, block_ref.kv_cache + ): + torch.testing.assert_close( + batched_state, ref_state, rtol=2e-4, atol=1e-3 + ) + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_block_batches_prefill_tokens_without_changing_results(): + config = _make_config() + vllm_config = VllmConfig(device_config=DeviceConfig("cpu")) + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1, backend="gloo") + try: + block_batched = RWKV7Block( + config=config, layer_idx=0, prefix="model.layers.0" + ) + _initialize_module_parameters(block_batched) + block_ref = RWKV7Block(config=config, layer_idx=0, prefix="model.layers.1") + block_ref.load_state_dict(block_batched.state_dict()) + + generator = torch.Generator().manual_seed(321) + state_shapes = block_batched.get_state_shape() + state_dtypes = block_batched.get_state_dtype() + + def make_cache() -> tuple[torch.Tensor, ...]: + return tuple( + torch.randn( + (2, *shape), + generator=generator, + dtype=dtype, + ) + for shape, dtype in zip(state_shapes, state_dtypes) + ) + + block_batched.kv_cache = make_cache() + block_ref.kv_cache = tuple( + cache.clone() for cache in block_batched.kv_cache + ) + + query_lens = [2, 3] + total_seq_lens = [2, 5] + state_indices = [0, 1] + hidden_states = torch.randn( + sum(query_lens), + config.hidden_size, + generator=generator, + dtype=torch.float32, + ) + metadata = _make_multi_prefill_metadata( + query_lens, + total_seq_lens, + state_indices, + device=torch.device("cpu"), + ) + + output_batched, v_first_batched = block_batched( + hidden_states, None, metadata + ) + + output_ref = torch.empty_like(hidden_states) + v_first_ref = torch.empty_like(hidden_states) + start = 0 + for slot_id, query_len, total_seq_len in zip( + state_indices, + query_lens, + total_seq_lens, + ): + end = start + query_len + states = block_ref._get_kv_state( + slot_id, + use_initial_state=total_seq_len > query_len, + ) + ( + out, + v_first_out, + attn_shift, + recurrent, + ffn_shift, + ) = block_ref._run_sequence( + hidden_states[start:end], + None, + *states, + ) + output_ref[start:end] = out + v_first_ref[start:end] = v_first_out + block_ref._store_kv_state(slot_id, attn_shift, recurrent, ffn_shift) + start = end + + torch.testing.assert_close(output_batched, output_ref) + torch.testing.assert_close(v_first_batched, v_first_ref) + for batched_state, ref_state in zip( + block_batched.kv_cache, block_ref.kv_cache + ): + torch.testing.assert_close(batched_state, ref_state) + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_mamba_state_copy_function_types(): + copy_funcs = RWKV7ForCausalLM.get_mamba_state_copy_func() + assert copy_funcs == ( + get_conv_copy_spec, + get_temporal_copy_spec, + get_conv_copy_spec, + ) + + +def test_rwkv7_uses_base_mamba_model_config(): + assert MODELS_CONFIG_MAP["RWKV7ForCausalLM"] is MambaModelConfig + + +def test_rwkv7_does_not_declare_mamba_prefix_caching_support(): + assert getattr(RWKV7ForCausalLM, "supports_mamba_prefix_caching", False) is False + + +def test_rwkv7_declares_bitsandbytes_weight_mapping_contract(): + # BitsAndBytes requires the model class to declare how checkpoint linear + # names map to packed vLLM modules. RWKV7 does not pack its projections. + assert RWKV7ForCausalLM.packed_modules_mapping == {} + + +def test_rwkv7_keeps_recurrent_state_update_projections_unquantized(): + class RecordingQuantConfig: + def __init__(self) -> None: + self.prefixes: list[str] = [] + + def get_quant_method(self, layer, prefix: str): + del layer + self.prefixes.append(prefix) + return UnquantizedLinearMethod() + + config = _make_config() + quant_config = RecordingQuantConfig() + vllm_config = VllmConfig(device_config=DeviceConfig("cpu")) + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1, backend="gloo") + try: + attention = RWKV7Attention( + config=config, + layer_idx=1, + quant_config=quant_config, + prefix="model.layers.1.attn", + ) + for projection in ( + attention.r_proj, + attention.o_proj, + ): + assert projection.quant_config is quant_config + for projection in (attention.k_proj, attention.v_proj): + assert projection.quant_config is None + for low_rank in ( + attention.w_lora, + attention.a_lora, + attention.v_lora, + attention.g_lora, + ): + assert low_rank.lora[0].quant_config is None + assert low_rank.lora[2].quant_config is None + assert set(quant_config.prefixes) == { + "model.layers.1.attn.r_proj", + "model.layers.1.attn.o_proj", + } + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_quantizes_untied_lm_head(tmp_path: Path): + class RecordingQuantConfig: + def __init__(self) -> None: + self.prefixes: list[str] = [] + + def get_quant_method(self, layer, prefix: str): + del layer + self.prefixes.append(prefix) + return UnquantizedLinearMethod() + + model_path = _write_test_model_config(tmp_path) + quant_config = RecordingQuantConfig() + vllm_config = _make_vllm_config(model_path, device="cpu") + vllm_config.quant_config = quant_config + + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1, backend="gloo") + try: + model = RWKV7ForCausalLM(vllm_config=vllm_config) + assert model.lm_head.quant_config is quant_config + assert "lm_head" in quant_config.prefixes + finally: + cleanup_dist_env_and_memory() + + +@pytest.mark.parametrize( + ("weight_dtype", "expected_hidden_dtype"), + [(torch.float32, torch.float32), (torch.int8, torch.float16)], +) +def test_rwkv7_compute_logits_preserves_activation_dtype_for_quantized_head( + weight_dtype: torch.dtype, expected_hidden_dtype: torch.dtype +): + lm_head = SimpleNamespace(weight=torch.empty(1, dtype=weight_dtype)) + model = SimpleNamespace( + lm_head=lm_head, + logits_processor=lambda _head, hidden_states: hidden_states, + ) + hidden_states = torch.ones(1, dtype=torch.float16) + + output = RWKV7ForCausalLM.compute_logits(model, hidden_states) + + assert output.dtype == expected_hidden_dtype + + +def test_rwkv7_prefix_caching_defaults_to_align(tmp_path: Path, monkeypatch): + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm_cache")) + + model_path = _write_test_model_config(tmp_path) + vllm_config = VllmConfig( + model_config=ModelConfig( + str(model_path), + trust_remote_code=False, + dtype="float32", + runner="generate", + ), + parallel_config=ParallelConfig( + tensor_parallel_size=1, + pipeline_parallel_size=1, + ), + cache_config=CacheConfig( + enable_prefix_caching=True, + mamba_cache_mode="none", + ), + device_config=DeviceConfig("cpu"), + ) + + assert vllm_config.model_config.supports_mamba_prefix_caching is False + assert vllm_config.cache_config.mamba_cache_mode == "align" + assert vllm_config.cache_config.mamba_block_size == ( + vllm_config.cache_config.block_size + ) + + +def test_rwkv7_prefix_caching_all_mode_falls_back_to_align(tmp_path: Path, monkeypatch): + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm_cache")) + + model_path = _write_test_model_config(tmp_path) + vllm_config = VllmConfig( + model_config=ModelConfig( + str(model_path), + trust_remote_code=False, + dtype="float32", + runner="generate", + ), + parallel_config=ParallelConfig( + tensor_parallel_size=1, + pipeline_parallel_size=1, + ), + cache_config=CacheConfig( + enable_prefix_caching=True, + mamba_cache_mode="all", + ), + device_config=DeviceConfig("cpu"), + ) + + assert vllm_config.model_config.supports_mamba_prefix_caching is False + assert vllm_config.cache_config.mamba_cache_mode == "align" + + +def test_rwkv7_block_uses_fp32_runtime_state_dtype(): + config = _make_config() + vllm_config = VllmConfig(device_config=DeviceConfig("cpu")) + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1, backend="gloo") + try: + block = RWKV7Block(config=config, layer_idx=0, prefix="model.layers.0") + assert block.get_state_dtype() == ( + torch.float32, + torch.float32, + torch.float32, + ) + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_model_keeps_weights_and_pp_buffers_in_model_dtype(tmp_path: Path): + model_path = _write_test_model_config(tmp_path) + vllm_config = _make_vllm_config( + model_path, + dtype="bfloat16", + device="cpu", + ) + + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="gloo", + ) + ensure_model_parallel_initialized(1, 1, backend="gloo") + try: + model = RWKV7ForCausalLM(vllm_config=vllm_config) + + parameter_dtypes = { + parameter.dtype + for parameter in model.parameters() + if parameter.is_floating_point() + } + assert parameter_dtypes == {torch.bfloat16} + assert model.model._pp_runtime_dtype() == torch.bfloat16 + + intermediate_tensors = model.make_empty_intermediate_tensors( + batch_size=2, + dtype=torch.float32, + device=torch.device("cpu"), + ) + assert intermediate_tensors["hidden_states"].dtype == torch.bfloat16 + assert intermediate_tensors["v_first"].dtype == torch.bfloat16 + + assert model.model.layers[0].get_state_dtype() == ( + torch.float32, + torch.float32, + torch.float32, + ) + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_reference_parity_full_forward(): + if pytest is None: + raise RuntimeError("pytest is required to run RWKV7 integration tests.") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for RWKV7 reference parity tests.") + + model_path, reference_cls = _require_reference_checkpoint() + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + inputs = tokenizer("Hello RWKV7, this is a parity check.", return_tensors="pt")[ + "input_ids" + ].to("cuda") + flat_input_ids = inputs[0] + positions = torch.arange(flat_input_ids.numel(), device="cuda", dtype=torch.long) + + reference_model = ( + reference_cls.from_pretrained(model_path, dtype=torch.float32).eval().to("cuda") + ) + vllm_config = _make_vllm_config(model_path) + + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="nccl", + ) + ensure_model_parallel_initialized(1, 1, backend="nccl") + try: + vllm_model = RWKV7ForCausalLM(vllm_config=vllm_config) + vllm_model.load_weights(reference_model.state_dict().items()) + vllm_model = vllm_model.eval().to("cuda", torch.float32) + + with torch.no_grad(): + reference_outputs = reference_model( + input_ids=inputs, + use_cache=False, + ) + reference_hidden = reference_model.model( + input_ids=inputs, + use_cache=False, + )[0][0] + + with torch.no_grad(), set_forward_context(None, vllm_config): + hidden_states = vllm_model( + input_ids=flat_input_ids, + positions=positions, + ) + logits = vllm_model.compute_logits(hidden_states) + + hidden_diff = (hidden_states - reference_hidden).abs() + logits_diff = (logits - reference_outputs.logits[0]).abs() + # The external FLA reference path is a test-only dependency and + # can show small (~1e-4) hidden-state drift versus vLLM on local + # checkpoints while still producing tightly matched logits. + assert hidden_diff.max().item() < 2e-4 + assert logits_diff.max().item() < 5e-5 + finally: + cleanup_dist_env_and_memory() + + +def test_rwkv7_reference_parity_prefill_decode(): + if pytest is None: + raise RuntimeError("pytest is required to run RWKV7 integration tests.") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for RWKV7 reference parity tests.") + + model_path, reference_cls = _require_reference_checkpoint() + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + prompt_ids = tokenizer("The capital of France is", return_tensors="pt")[ + "input_ids" + ].to("cuda") + + reference_model = ( + reference_cls.from_pretrained(model_path, dtype=torch.float32).eval().to("cuda") + ) + vllm_config = _make_vllm_config(model_path) + + with set_current_vllm_config(vllm_config): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=_local_distributed_init_method(), + backend="nccl", + ) + ensure_model_parallel_initialized(1, 1, backend="nccl") + try: + vllm_model = RWKV7ForCausalLM(vllm_config=vllm_config) + vllm_model.load_weights(reference_model.state_dict().items()) + vllm_model = vllm_model.eval().to("cuda", torch.float32) + _allocate_kv_cache(vllm_model, device=torch.device("cuda")) + + prompt_flat = prompt_ids[0] + prompt_positions = torch.arange( + prompt_flat.numel(), device="cuda", dtype=torch.long + ) + prompt_metadata = { + layer.prefix: _make_prefill_metadata( + prompt_ids.shape[1], device=torch.device("cuda") + ) + for layer in vllm_model.model.layers + } + + with torch.no_grad(): + reference_prompt_logits = reference_model( + input_ids=prompt_ids, use_cache=False + ).logits[0] + + with torch.no_grad(), set_forward_context(prompt_metadata, vllm_config): + hidden_states = vllm_model( + input_ids=prompt_flat, + positions=prompt_positions, + ) + logits = vllm_model.compute_logits(hidden_states) + + assert (logits - reference_prompt_logits).abs().max().item() < 5e-5 + + next_token = logits[-1].argmax().view(1) + reference_first_token = reference_prompt_logits[-1].argmax().view(1) + assert int(next_token.item()) == int(reference_first_token.item()) + + generated_vllm = [int(next_token.item())] + generated_ref = [int(reference_first_token.item())] + current_ids = prompt_ids.clone() + + for _ in range(3): + total_seq_len = current_ids.shape[1] + 1 + decode_metadata = { + layer.prefix: _make_decode_metadata( + total_seq_len, device=torch.device("cuda") + ) + for layer in vllm_model.model.layers + } + position = torch.tensor( + [current_ids.shape[1]], device="cuda", dtype=torch.long + ) + + with torch.no_grad(), set_forward_context(decode_metadata, vllm_config): + hidden_states = vllm_model( + input_ids=next_token, + positions=position, + ) + logits = vllm_model.compute_logits(hidden_states) + + full_ids = torch.cat([current_ids, next_token.view(1, 1)], dim=1) + with torch.no_grad(): + reference_last_logits = reference_model( + input_ids=full_ids, + use_cache=False, + ).logits[0, -1] + + assert (logits[-1] - reference_last_logits).abs().max().item() < 5e-5 + + next_token = logits[-1].argmax().view(1) + reference_next_token = reference_last_logits.argmax().view(1) + assert int(next_token.item()) == int(reference_next_token.item()) + + generated_vllm.append(int(next_token.item())) + generated_ref.append(int(reference_next_token.item())) + current_ids = full_ids + + assert generated_vllm == generated_ref + assert tokenizer.decode(generated_vllm) == tokenizer.decode(generated_ref) + finally: + cleanup_dist_env_and_memory() diff --git a/tests/models/registry.py b/tests/models/registry.py index 00a84783bf11..1008ea24ddb8 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -402,6 +402,10 @@ def check_available_online( "random": "yujiepan/mamba2-codestral-v0.1-tiny-random", }, ), + "RWKV7ForCausalLM": _HfExamplesInfo( + "RWKV/RWKV7-Goose-World2.8-0.1B-HF", + trust_remote_code=True, + ), "FalconMambaForCausalLM": _HfExamplesInfo("tiiuae/falcon-mamba-7b-instruct"), "MiniCPMForCausalLM": _HfExamplesInfo( "openbmb/MiniCPM-2B-sft-bf16", trust_remote_code=True diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index d977cff96d02..b7d0abf47035 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -9,13 +9,23 @@ _test_online_quant_peak_mem_impl, is_quant_method_supported, ) -from vllm.model_executor.layers.linear import UnquantizedLinearMethod +from vllm.config.quantization import QuantizationConfigArgs +from vllm.model_executor.layers.linear import ( + ReplicatedLinear, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, +) from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerBlockOnlineLinearMethod, Fp8PerBlockOnlineMoEMethod, Fp8PerTensorOnlineLinearMethod, Fp8PerTensorOnlineMoEMethod, ) +from vllm.model_executor.layers.quantization.online.int8 import ( + Int8OnlineLinearMethod, +) from vllm.model_executor.layers.quantization.online.nvfp4 import ( Nvfp4OnlineMoEMethod, ) @@ -23,6 +33,54 @@ from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe +def test_online_int8_linear_quantizes_per_output_channel() -> None: + class RecordingKernel: + def __init__(self) -> None: + self.processed = 0 + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self.processed += 1 + assert layer.weight.dtype == torch.int8 + + layer = torch.nn.Module() + original = torch.tensor( + [[-2.0, -0.25, 1.0, 2.0], [-0.1, 0.0, 0.05, 0.1]], + dtype=torch.float16, + ) + layer.register_parameter( + "weight", torch.nn.Parameter(original.clone(), requires_grad=False) + ) + + method = Int8OnlineLinearMethod() + method.kernel = RecordingKernel() + method.process_weights_after_loading(layer) + method.process_weights_after_loading(layer) + + assert method.kernel.processed == 1 + assert layer.weight_scale.shape == (2, 1) + reconstructed = layer.weight.float() * layer.weight_scale + assert torch.allclose(reconstructed, original.float(), atol=0.01, rtol=0) + assert layer.input_scale is None + assert layer.input_zero_point is None + assert layer.azp_adj is None + + +def test_online_int8_dispatches_dense_linear_and_lm_head() -> None: + from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead + + config = OnlineQuantizationConfig( + QuantizationConfigArgs(linear="int8_per_channel_static") + ) + layers = [ + object.__new__(ReplicatedLinear), + object.__new__(ParallelLMHead), + ] + + for layer in layers: + method = config.get_quant_method(layer, "linear") + assert isinstance(method, Int8OnlineLinearMethod) + + @pytest.mark.skipif( not is_quant_method_supported("fp8"), reason="FP8 is not supported on this GPU type.", diff --git a/tests/quantization/test_torchao.py b/tests/quantization/test_torchao.py index a724803b9a18..f0aea6ff4725 100644 --- a/tests/quantization/test_torchao.py +++ b/tests/quantization/test_torchao.py @@ -27,6 +27,23 @@ def on_gfx950() -> bool: ) >= version.parse("0.18.0") +@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available") +def test_torchao_quantizes_parallel_lm_head(): + from torchao.quantization import Int8WeightOnlyConfig + + from vllm.model_executor.layers.quantization.torchao import ( + TorchAOConfig, + TorchAOLinearMethod, + ) + from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead + + # get_quant_method only dispatches on the module type, so construction of + # the sharded weights and a distributed environment is unnecessary here. + lm_head = object.__new__(ParallelLMHead) + method = TorchAOConfig(Int8WeightOnlyConfig()).get_quant_method(lm_head, "lm_head") + assert isinstance(method, TorchAOLinearMethod) + + @pytest.mark.skipif( current_platform.is_rocm() and current_platform.is_fp8_fnuz(), reason="Only fp8_fnuz supported on CDNA3 architecture", diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index fbc5dcff40e8..d7680f195ecf 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -772,6 +772,7 @@ class CompilationConfig: "vllm::gdn_attention_core_xpu", "vllm::olmo_hybrid_gdn_full_forward", "vllm::kda_attention", + "vllm::rwkv7_block_forward", "vllm::sparse_attn_indexer", "vllm::rocm_aiter_sparse_attn_indexer", "vllm::deepseek_v4_attention", diff --git a/vllm/envs.py b/vllm/envs.py index fb54619c7485..8735efc9ee44 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -122,6 +122,7 @@ VLLM_SKIP_P2P_CHECK: bool = False VLLM_DISABLED_KERNELS: list[str] = [] VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE: bool = True + VLLM_RWKV7_KERNEL: Literal["auto", "torch", "triton"] = "auto" VLLM_DISABLE_PYNCCL: bool = False VLLM_USE_OINK_OPS: bool = False VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True @@ -1164,6 +1165,15 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE": lambda: bool( int(os.getenv("VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE", "1")) ), + # Select the RWKV7 recurrent implementation. ``auto`` is intentionally + # fail-closed to the PyTorch reference path until the Triton kernel passes + # the long-horizon greedy-token parity gate. + "VLLM_RWKV7_KERNEL": env_with_choices( + "VLLM_RWKV7_KERNEL", + "auto", + ["auto", "torch", "triton"], + case_sensitive=False, + ), # Disable pynccl (using torch.distributed instead) "VLLM_DISABLE_PYNCCL": lambda: ( os.getenv("VLLM_DISABLE_PYNCCL", "False").lower() in ("true", "1") diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index 3f8a741d357d..4b976f5b4369 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -34,6 +34,7 @@ Fp8PtpcOnlineMoEMethod, ) from vllm.model_executor.layers.quantization.online.int8 import ( + Int8OnlineLinearMethod, Int8OnlineMoEMethod, ) from vllm.model_executor.layers.quantization.online.mxfp8 import ( @@ -64,6 +65,7 @@ kFp8Static128BlockSym: Fp8PerBlockOnlineLinearMethod, kFp8StaticChannelSym: Fp8PtpcOnlineLinearMethod, kMxfp8Dynamic: Mxfp8OnlineLinearMethod, + kInt8StaticChannelSym: Int8OnlineLinearMethod, } _ONLINE_MOE_METHODS: dict[QuantKey, type] = { @@ -150,7 +152,11 @@ def _dispatch( def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": - if isinstance(layer, LinearBase): + from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + ) + + if isinstance(layer, (LinearBase, ParallelLMHead)): if should_ignore_layer( prefix, ignore=self.ignored_layers, diff --git a/vllm/model_executor/layers/quantization/online/int8.py b/vllm/model_executor/layers/quantization/online/int8.py index ef76594164a8..5220172bb0b2 100644 --- a/vllm/model_executor/layers/quantization/online/int8.py +++ b/vllm/model_executor/layers/quantization/online/int8.py @@ -5,12 +5,14 @@ import torch from torch.nn import Module +from torch.nn.parameter import Parameter if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) +from vllm.model_executor.kernels.linear import init_int8_linear_kernel from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.fused_moe.oracle.int8 import ( convert_to_int8_moe_kernel_format, @@ -18,6 +20,7 @@ make_int8_moe_quant_config, select_int8_moe_backend, ) +from vllm.model_executor.layers.linear import LinearMethodBase from vllm.model_executor.layers.quantization.online.moe_base import ( OnlineMoEMethodBase, ) @@ -25,7 +28,67 @@ kInt8DynamicTokenSym, kInt8StaticChannelSym, ) -from vllm.model_executor.utils import replace_parameter +from vllm.model_executor.utils import replace_parameter, set_weight_attrs + + +class Int8OnlineLinearMethod(LinearMethodBase): + """Online per-channel INT8 weights with dynamic per-token activations.""" + + def create_weights( + self, + layer: Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del input_size, output_size + layer.logical_widths = output_partition_sizes + weight = Parameter( + torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + self.kernel = init_int8_linear_kernel( + is_channelwise=True, + is_static_input_scheme=False, + input_symmetric=True, + module_name=self.__class__.__name__, + ) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + weight = layer.weight.float() + weight_scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127 + weight = (weight / weight_scale).round().clamp(-127, 127).to(torch.int8) + + replace_parameter(layer, "weight", weight) + layer.register_parameter( + "weight_scale", Parameter(weight_scale, requires_grad=False) + ) + layer.register_parameter("input_scale", None) + layer.register_parameter("input_zero_point", None) + layer.register_parameter("azp_adj", None) + self.kernel.process_weights_after_loading(layer) + layer._already_called_process_weights_after_loading = True + + def apply( + self, + layer: Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.kernel.apply_weights(layer, x, bias) class Int8OnlineMoEMethod(OnlineMoEMethodBase): diff --git a/vllm/model_executor/layers/quantization/torchao.py b/vllm/model_executor/layers/quantization/torchao.py index 15399cfd39b4..a0c8578716a8 100644 --- a/vllm/model_executor/layers/quantization/torchao.py +++ b/vllm/model_executor/layers/quantization/torchao.py @@ -242,7 +242,15 @@ def from_config_dict_json(cls, config_dict_json: str) -> "TorchAOConfig": def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": - if not isinstance(layer, LinearBase): + # ParallelLMHead is implemented on top of VocabParallelEmbedding, but + # its sampler path is a linear projection and accepts LinearMethodBase + # implementations. Treat it like the other linear layers so TorchAO + # can quantize large untied output heads. + from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + ) + + if not isinstance(layer, (LinearBase, ParallelLMHead)): return None from torchao.quantization import ModuleFqnToConfig diff --git a/vllm/model_executor/layers/rwkv7.py b/vllm/model_executor/layers/rwkv7.py new file mode 100644 index 000000000000..2e4f50a83172 --- /dev/null +++ b/vllm/model_executor/layers/rwkv7.py @@ -0,0 +1,471 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RWKV7 recurrent state operators. + +The runtime stores state as ``[batch, heads, value_dim, key_dim]``. This is the +transpose of the equation's key-by-value matrix and lets each Triton program +own one value row without cross-program reductions. +""" + +import os +from dataclasses import dataclass +from typing import Literal, cast + +import torch + +import vllm.envs as envs +from vllm.triton_utils import HAS_TRITON, tl, triton + +RWKV7KernelBackend = Literal["auto", "torch", "triton"] + +# Do not promote Triton through ``auto`` until long-horizon greedy generation +# is exact. Users can still opt in explicitly for benchmarking and validation. +_TRITON_AUTO_ENABLED = False + +# A multi-warp reduction is faster in isolation but changes fp32 accumulation +# enough to flip near-tied greedy logits after recurrent state accumulation. +_TRITON_NUM_WARPS = 1 + + +@dataclass(frozen=True) +class RWKV7KernelParityReport: + """Numerical comparison between the Torch and Triton recurrent paths.""" + + output_close: bool + state_close: bool + first_output_mismatch_token: int | None + first_state_mismatch_sequence: int | None + output_max_abs_error: float + output_max_rel_error: float + output_cosine_similarity: float + state_max_abs_error: float + state_max_rel_error: float + state_cosine_similarity: float + atol: float + rtol: float + + +def resolve_rwkv7_kernel_backend( + requested: RWKV7KernelBackend | None = None, + *, + input_tensor: torch.Tensor | None = None, + state: torch.Tensor | None = None, +) -> Literal["torch", "triton"]: + """Resolve and validate the recurrent backend. + + ``auto`` deliberately resolves to ``torch`` while the fused kernel is + experimental. Explicit ``triton`` requests fail instead of silently + falling back, which keeps benchmarks and correctness gates honest. + """ + fallback = cast( + RWKV7KernelBackend, + os.getenv("VLLM_RWKV7_KERNEL", "auto"), + ) + backend: str = ( + requested + if requested is not None + else getattr(envs, "VLLM_RWKV7_KERNEL", fallback) + ) + backend = backend.lower() + if backend not in ("auto", "torch", "triton"): + raise ValueError( + f"Invalid RWKV7 kernel backend {backend!r}; expected auto, torch, or triton" + ) + if backend == "auto": + backend = "triton" if _TRITON_AUTO_ENABLED else "torch" + if backend == "triton": + if not HAS_TRITON: + raise RuntimeError( + "RWKV7 Triton kernel requested, but Triton is unavailable" + ) + if input_tensor is not None and not input_tensor.is_cuda: + raise RuntimeError("RWKV7 Triton kernel requires CUDA or ROCm tensors") + if state is not None and state.dtype != torch.float32: + raise RuntimeError("RWKV7 Triton kernel requires fp32 recurrent state") + return cast(Literal["torch", "triton"], backend) + + +def _rwkv7_step_reference( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + state: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + state_dot_kk = (state * kk.unsqueeze(-2)).sum(dim=-1) + new_state = ( + state * torch.exp(w).unsqueeze(-2) + + v.unsqueeze(-1) * k.unsqueeze(-2) + - state_dot_kk.unsqueeze(-1) * (kk * a).unsqueeze(-2) + ) + output = (new_state * r.unsqueeze(-2)).sum(dim=-1) + return output, new_state + + +def _rwkv7_scan_packed_reference( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + state: torch.Tensor, + query_start_loc: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + output = torch.empty_like(v) + final_state = torch.empty_like(state) + for seq_idx in range(query_start_loc.numel() - 1): + start = int(query_start_loc[seq_idx].item()) + end = int(query_start_loc[seq_idx + 1].item()) + seq_state = state[seq_idx : seq_idx + 1] + for token_idx in range(start, end): + token_output, seq_state = _rwkv7_step_reference( + r[token_idx : token_idx + 1], + w[token_idx : token_idx + 1], + k[token_idx : token_idx + 1], + v[token_idx : token_idx + 1], + kk[token_idx : token_idx + 1], + a[token_idx : token_idx + 1], + seq_state, + ) + output[token_idx].copy_(token_output[0]) + final_state[seq_idx].copy_(seq_state[0]) + return output, final_state + + +if HAS_TRITON: + + @triton.jit + def _rwkv7_recurrent_step_kernel( + r_ptr, + w_ptr, + k_ptr, + v_ptr, + kk_ptr, + a_ptr, + state_ptr, + output_ptr, + final_state_ptr, + K: tl.constexpr, + V: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + row_id = tl.program_id(0) + value_idx = row_id % V + batch_head_idx = row_id // V + + offsets = tl.arange(0, BLOCK_K) + mask = offsets < K + vector_base = batch_head_idx * K + state_base = row_id * K + + state = tl.load(state_ptr + state_base + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + r = tl.load(r_ptr + vector_base + offsets, mask=mask, other=0.0).to(tl.float32) + w = tl.load(w_ptr + vector_base + offsets, mask=mask, other=0.0).to(tl.float32) + k = tl.load(k_ptr + vector_base + offsets, mask=mask, other=0.0).to(tl.float32) + kk = tl.load(kk_ptr + vector_base + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + a = tl.load(a_ptr + vector_base + offsets, mask=mask, other=0.0).to(tl.float32) + value = tl.load(v_ptr + batch_head_idx * V + value_idx).to(tl.float32) + + state_dot_kk = tl.sum(state * kk, axis=0) + state = state * tl.exp(w) + value * k - state_dot_kk * kk * a + output = tl.sum(state * r, axis=0) + + tl.store(final_state_ptr + state_base + offsets, state, mask=mask) + tl.store(output_ptr + batch_head_idx * V + value_idx, output) + + @triton.jit + def _rwkv7_recurrent_scan_packed_kernel( + r_ptr, + w_ptr, + k_ptr, + v_ptr, + kk_ptr, + a_ptr, + state_ptr, + query_start_loc_ptr, + output_ptr, + final_state_ptr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + row_id = tl.program_id(0) + value_idx = row_id % V + seq_head_idx = row_id // V + head_idx = seq_head_idx % H + seq_idx = seq_head_idx // H + + offsets = tl.arange(0, BLOCK_K) + mask = offsets < K + state_base = row_id * K + state = tl.load(state_ptr + state_base + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + + token_idx = tl.load(query_start_loc_ptr + seq_idx).to(tl.int64) + end = tl.load(query_start_loc_ptr + seq_idx + 1).to(tl.int64) + while token_idx < end: + vector_base = (token_idx * H + head_idx) * K + value_base = (token_idx * H + head_idx) * V + r = tl.load(r_ptr + vector_base + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + w = tl.load(w_ptr + vector_base + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + k = tl.load(k_ptr + vector_base + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + kk = tl.load(kk_ptr + vector_base + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + a = tl.load(a_ptr + vector_base + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + value = tl.load(v_ptr + value_base + value_idx).to(tl.float32) + + state_dot_kk = tl.sum(state * kk, axis=0) + state = state * tl.exp(w) + value * k - state_dot_kk * kk * a + output = tl.sum(state * r, axis=0) + tl.store(output_ptr + value_base + value_idx, output) + token_idx += 1 + + tl.store(final_state_ptr + state_base + offsets, state, mask=mask) + + +def _validate_inputs( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + state: torch.Tensor, +) -> tuple[int, int, int]: + if r.ndim != 3: + raise ValueError("r/w/k/kk/a must be shaped [batch, heads, key_dim]") + if any(t.shape != r.shape for t in (w, k, kk, a)): + raise ValueError("r/w/k/kk/a must have identical shapes") + if v.ndim != 3 or v.shape[:2] != r.shape[:2]: + raise ValueError("v must be shaped [batch, heads, value_dim]") + batch, heads, key_dim = r.shape + value_dim = v.shape[-1] + if state.shape != (batch, heads, value_dim, key_dim): + raise ValueError("state must be shaped [batch, heads, value_dim, key_dim]") + if any(t.device != r.device for t in (w, k, v, kk, a, state)): + raise ValueError("all RWKV7 recurrent inputs must be on the same device") + return heads, key_dim, value_dim + + +def rwkv7_recurrent_step( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + state: torch.Tensor, + *, + backend: RWKV7KernelBackend | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run one batched RWKV7 recurrent update without modifying ``state``.""" + heads, key_dim, value_dim = _validate_inputs(r, w, k, v, kk, a, state) + selected_backend = resolve_rwkv7_kernel_backend( + backend, input_tensor=r, state=state + ) + if selected_backend == "torch": + return _rwkv7_step_reference(r, w, k, v, kk, a, state) + + output = torch.empty_like(v) + final_state = torch.empty_like(state) + block_key = triton.next_power_of_2(key_dim) + _rwkv7_recurrent_step_kernel[(r.shape[0] * heads * value_dim,)]( + r.contiguous(), + w.contiguous(), + k.contiguous(), + v.contiguous(), + kk.contiguous(), + a.contiguous(), + state.contiguous(), + output, + final_state, + K=key_dim, + V=value_dim, + BLOCK_K=block_key, + num_warps=_TRITON_NUM_WARPS, + ) + return output, final_state + + +def rwkv7_recurrent_scan_packed( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + state: torch.Tensor, + query_start_loc: torch.Tensor, + *, + backend: RWKV7KernelBackend | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run an RWKV7 scan over sequences delimited by ``query_start_loc``.""" + if query_start_loc.ndim != 1 or query_start_loc.numel() < 2: + raise ValueError("query_start_loc must contain one start per sequence plus end") + num_sequences = query_start_loc.numel() - 1 + if r.ndim != 3: + raise ValueError("packed r/w/k/kk/a must be shaped [tokens, heads, key_dim]") + if state.shape[0] != num_sequences: + raise ValueError("state batch size must match query_start_loc") + if any(t.shape != r.shape for t in (w, k, kk, a)): + raise ValueError("packed r/w/k/kk/a must have identical shapes") + if v.ndim != 3 or v.shape[:2] != r.shape[:2]: + raise ValueError("packed v must match the token and head dimensions") + _, heads, key_dim = r.shape + value_dim = v.shape[-1] + if state.shape != (num_sequences, heads, value_dim, key_dim): + raise ValueError("state must be shaped [sequences, heads, value_dim, key_dim]") + if any(t.device != r.device for t in (w, k, v, kk, a, state)): + raise ValueError("all RWKV7 recurrent inputs must be on the same device") + if query_start_loc.device != r.device: + raise ValueError("query_start_loc must be on the recurrent input device") + + selected_backend = resolve_rwkv7_kernel_backend( + backend, input_tensor=r, state=state + ) + if selected_backend == "torch": + return _rwkv7_scan_packed_reference(r, w, k, v, kk, a, state, query_start_loc) + + output = torch.empty_like(v) + final_state = torch.empty_like(state) + block_key = triton.next_power_of_2(key_dim) + _rwkv7_recurrent_scan_packed_kernel[(num_sequences * heads * value_dim,)]( + r.contiguous(), + w.contiguous(), + k.contiguous(), + v.contiguous(), + kk.contiguous(), + a.contiguous(), + state.contiguous(), + query_start_loc.contiguous(), + output, + final_state, + H=heads, + K=key_dim, + V=value_dim, + BLOCK_K=block_key, + num_warps=_TRITON_NUM_WARPS, + ) + return output, final_state + + +def _error_metrics( + actual: torch.Tensor, expected: torch.Tensor +) -> tuple[float, float, float]: + if actual.numel() == 0: + return 0.0, 0.0, 1.0 + actual_fp32 = actual.to(torch.float32).flatten() + expected_fp32 = expected.to(torch.float32).flatten() + error = (actual_fp32 - expected_fp32).abs() + max_abs_error = error.max().item() + max_rel_error = (error / expected_fp32.abs().clamp_min(1e-12)).max().item() + denominator = torch.linalg.vector_norm(actual_fp32) * torch.linalg.vector_norm( + expected_fp32 + ) + if denominator.item() == 0.0: + cosine_similarity = float(torch.equal(actual_fp32, expected_fp32)) + else: + cosine_similarity = ( + torch.dot(actual_fp32, expected_fp32).div(denominator).item() + ) + return max_abs_error, max_rel_error, cosine_similarity + + +def _first_mismatch( + actual: torch.Tensor, + expected: torch.Tensor, + *, + atol: float, + rtol: float, +) -> int | None: + if actual.shape[0] == 0: + return None + mismatch = (actual - expected).abs() > atol + rtol * expected.abs() + mismatch_by_item = mismatch.reshape(actual.shape[0], -1).any(dim=1) + indices = mismatch_by_item.nonzero() + return None if indices.numel() == 0 else int(indices[0].item()) + + +@torch.no_grad() +def diagnose_rwkv7_recurrent_scan_packed( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + state: torch.Tensor, + query_start_loc: torch.Tensor, + *, + atol: float = 1e-5, + rtol: float = 1e-5, +) -> RWKV7KernelParityReport: + """Compare packed Torch and Triton scans and locate the first mismatch. + + This diagnostic intentionally runs both implementations and synchronizes + scalar metrics back to the host. It is for correctness qualification, not + serving. ``first_output_mismatch_token`` indexes the packed token axis; + ``first_state_mismatch_sequence`` indexes the final-state batch axis. + """ + expected_output, expected_state = rwkv7_recurrent_scan_packed( + r, + w, + k, + v, + kk, + a, + state, + query_start_loc, + backend="torch", + ) + actual_output, actual_state = rwkv7_recurrent_scan_packed( + r, + w, + k, + v, + kk, + a, + state, + query_start_loc, + backend="triton", + ) + output_metrics = _error_metrics(actual_output, expected_output) + state_metrics = _error_metrics(actual_state, expected_state) + first_output_mismatch = _first_mismatch( + actual_output, expected_output, atol=atol, rtol=rtol + ) + first_state_mismatch = _first_mismatch( + actual_state, expected_state, atol=atol, rtol=rtol + ) + return RWKV7KernelParityReport( + output_close=first_output_mismatch is None, + state_close=first_state_mismatch is None, + first_output_mismatch_token=first_output_mismatch, + first_state_mismatch_sequence=first_state_mismatch, + output_max_abs_error=output_metrics[0], + output_max_rel_error=output_metrics[1], + output_cosine_similarity=output_metrics[2], + state_max_abs_error=state_metrics[0], + state_max_rel_error=state_metrics[1], + state_cosine_similarity=state_metrics[2], + atol=atol, + rtol=rtol, + ) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 3219a5edcdac..d3a7e65ad794 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -886,6 +886,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig, "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, + "RWKV7ForCausalLM": MambaModelConfig, "UnlimitedOCRForCausalLM": UnlimitedOCRForCausalLMConfig, "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, "XLMRobertaModel": JinaRobertaModelConfig, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 697f55b37270..002a310657b8 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -195,6 +195,7 @@ "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"), "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"), "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"), + "RWKV7ForCausalLM": ("rwkv7", "RWKV7ForCausalLM"), "RWForCausalLM": ("falcon", "FalconForCausalLM"), "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"), "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"), diff --git a/vllm/model_executor/models/rwkv7.py b/vllm/model_executor/models/rwkv7.py new file mode 100644 index 000000000000..d9f4342384c7 --- /dev/null +++ b/vllm/model_executor/models/rwkv7.py @@ -0,0 +1,1539 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only RWKV7 model.""" + +from collections.abc import Iterable +from itertools import islice + +import torch +import torch.nn.functional as F +from torch import nn +from transformers.activations import ACT2FN as HF_ACT2FN + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config +from vllm.distributed.parallel_state import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + model_parallel_is_initialized, +) +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.abstract import MambaBase +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rwkv7 import ( + rwkv7_recurrent_scan_packed, + rwkv7_recurrent_step, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import ( + HasInnerState, + IsAttentionFree, + SupportsPP, +) +from vllm.sequence import IntermediateTensors +from vllm.utils.torch_utils import direct_register_custom_op +from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum + +from .utils import AutoWeightsLoader, PPMissingLayer, make_layers, maybe_prefix + +LOG_DECAY_SCALE = -0.6065306597126334 +RWKV7_RUNTIME_DTYPE = torch.float32 + + +def get_tp_world_size() -> int: + return ( + get_tensor_model_parallel_world_size() if model_parallel_is_initialized() else 1 + ) + + +def get_tp_rank() -> int: + return get_tensor_model_parallel_rank() if model_parallel_is_initialized() else 0 + + +def sqrelu(x: torch.Tensor) -> torch.Tensor: + return torch.relu(x).square() + + +def get_activation_fn(name: str): + if name == "sqrelu": + return sqrelu + if name not in HF_ACT2FN: + raise ValueError(f"Unsupported RWKV7 activation: {name}") + return HF_ACT2FN[name] + + +def token_shift_with_cache( + hidden_states: torch.Tensor, cached_state: torch.Tensor | None +) -> tuple[torch.Tensor, torch.Tensor]: + delta = torch.empty_like(hidden_states) + if hidden_states.shape[0] == 0: + final_state = ( + cached_state if cached_state is not None else hidden_states.new_empty(0) + ) + return delta, final_state + + if cached_state is None: + delta[0] = -hidden_states[0] + else: + delta[0] = cached_state.to(hidden_states.dtype) - hidden_states[0] + if hidden_states.shape[0] > 1: + delta[1:] = hidden_states[:-1] - hidden_states[1:] + final_state = hidden_states[-1].to( + cached_state.dtype if cached_state is not None else hidden_states.dtype + ) + return delta, final_state + + +def token_shift_with_cache_varlen( + hidden_states: torch.Tensor, + query_start_loc: torch.Tensor, + cached_state: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + delta = torch.empty_like(hidden_states) + if hidden_states.shape[0] == 0: + if cached_state is not None: + return delta, cached_state + return delta, hidden_states.new_empty((0, hidden_states.shape[-1])) + + delta[1:] = hidden_states[:-1] - hidden_states[1:] + first_token_indices = query_start_loc[:-1].to(dtype=torch.long) + if cached_state is None: + delta[first_token_indices] = -hidden_states.index_select(0, first_token_indices) + else: + delta[first_token_indices] = cached_state.to( + hidden_states.dtype + ) - hidden_states.index_select(0, first_token_indices) + + last_token_indices = (query_start_loc[1:] - 1).to(dtype=torch.long) + final_state = hidden_states.index_select(0, last_token_indices).to( + cached_state.dtype if cached_state is not None else hidden_states.dtype + ) + return delta, final_state + + +def _custom_op_optional_tensor( + tensor: torch.Tensor | None, + *, + like: torch.Tensor, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + if tensor is not None: + return tensor + return like.new_empty((0,), dtype=dtype or like.dtype) + + +def _custom_op_tensor_or_none(tensor: torch.Tensor) -> torch.Tensor | None: + return None if tensor.numel() == 0 else tensor + + +def _rwkv7_recurrent_step( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + recurrent_state: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + unbatched = recurrent_state.ndim == 3 + if unbatched: + r, w, k, v, kk, a, recurrent_state = ( + tensor.unsqueeze(0) for tensor in (r, w, k, v, kk, a, recurrent_state) + ) + output, final_state = rwkv7_recurrent_step( + r, + w, + k, + v, + kk, + a, + recurrent_state, + ) + if unbatched: + return output[0], final_state[0] + return output, final_state + + +def _rwkv7_recurrent_scan( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + initial_state: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + if initial_state is None: + recurrent_state = torch.zeros( + r.shape[1], + v.shape[2], + r.shape[2], + device=r.device, + dtype=torch.float32, + ) + else: + recurrent_state = initial_state.to(torch.float32) + query_start_loc = torch.tensor([0, r.shape[0]], dtype=torch.int32, device=r.device) + output, final_state = rwkv7_recurrent_scan_packed( + r, + w, + k, + v, + kk, + a, + recurrent_state.unsqueeze(0), + query_start_loc, + ) + return output, final_state[0] + + +def _rwkv7_recurrent_scan_varlen( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kk: torch.Tensor, + a: torch.Tensor, + query_start_loc: torch.Tensor, + initial_state: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + num_sequences = query_start_loc.numel() - 1 + if initial_state is None: + initial_state = torch.zeros( + num_sequences, + r.shape[1], + v.shape[2], + r.shape[2], + device=r.device, + dtype=torch.float32, + ) + return rwkv7_recurrent_scan_packed( + r, + w, + k, + v, + kk, + a, + initial_state.to(torch.float32), + query_start_loc, + ) + + +def rwkv7_attention( + hidden_states: torch.Tensor, + cached_shift_state: torch.Tensor, + recurrent_state: torch.Tensor, + v_first: torch.Tensor, + output: torch.Tensor, + final_shift_state: torch.Tensor, + final_recurrent_state: torch.Tensor, + v_first_out: torch.Tensor, + layer_name: str, +) -> None: + forward_context = get_forward_context() + self = forward_context.no_compile_layers[layer_name] + out, shift_state, recurrent, first_value = self._forward( + hidden_states, + _custom_op_tensor_or_none(cached_shift_state), + _custom_op_tensor_or_none(recurrent_state), + _custom_op_tensor_or_none(v_first), + ) + output.copy_(out) + final_shift_state.copy_(shift_state.to(final_shift_state.dtype)) + final_recurrent_state.copy_(recurrent.to(final_recurrent_state.dtype)) + v_first_out.copy_(first_value.to(v_first_out.dtype)) + + +def rwkv7_attention_fake( + hidden_states: torch.Tensor, + cached_shift_state: torch.Tensor, + recurrent_state: torch.Tensor, + v_first: torch.Tensor, + output: torch.Tensor, + final_shift_state: torch.Tensor, + final_recurrent_state: torch.Tensor, + v_first_out: torch.Tensor, + layer_name: str, +) -> None: + return + + +direct_register_custom_op( + op_name="rwkv7_attention", + op_func=rwkv7_attention, + mutates_args=[ + "output", + "final_shift_state", + "final_recurrent_state", + "v_first_out", + ], + fake_impl=rwkv7_attention_fake, + # These wrappers are only exercised on CUDA tensors. Register them + # explicitly on the CUDA dispatch key so they remain available even when + # current_platform resolves to an unspecified/CPU platform in test envs. + dispatch_key="CUDA", +) + + +def rwkv7_block_forward( + hidden_states: torch.Tensor, + v_first: torch.Tensor, + output: torch.Tensor, + v_first_out: torch.Tensor, + layer_name: str, +) -> None: + forward_context = get_forward_context() + self = forward_context.no_compile_layers[layer_name] + self._forward_runtime( + hidden_states, + _custom_op_tensor_or_none(v_first), + output, + v_first_out, + ) + + +def rwkv7_block_forward_fake( + hidden_states: torch.Tensor, + v_first: torch.Tensor, + output: torch.Tensor, + v_first_out: torch.Tensor, + layer_name: str, +) -> None: + return + + +direct_register_custom_op( + op_name="rwkv7_block_forward", + op_func=rwkv7_block_forward, + mutates_args=["output", "v_first_out"], + fake_impl=rwkv7_block_forward_fake, + dispatch_key="CUDA", +) + + +class RWKV7LoRA(nn.Module): + def __init__( + self, + input_dim: int, + output_dim: int, + low_rank_dim: int, + bias: bool, + activation: str | None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + # These low-rank projections generate decay, gate, and interpolation + # parameters. Quantizing them saves little memory, is numerically + # sensitive, and dispatches poorly on common weight-only kernels when + # the rank is 32/64/128. Keep them in the model dtype while the large + # recurrent and FFN projections use ``quant_config``. + del quant_config + if activation is None: + act = nn.Identity() + elif activation == "sigmoid": + act = nn.Sigmoid() + elif activation == "tanh": + act = nn.Tanh() + elif activation == "relu": + act = nn.ReLU() + else: + raise ValueError(f"Unsupported RWKV7 LoRA activation: {activation}") + + self.lora = nn.Sequential( + ReplicatedLinear( + input_dim, + low_rank_dim, + bias=False, + quant_config=None, + prefix=f"{prefix}.lora.0", + ), + act, + ColumnParallelLinear( + low_rank_dim, + output_dim, + bias=bias, + quant_config=None, + prefix=f"{prefix}.lora.2", + ), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.lora[0](x) + x = self.lora[1](x) + x, bias = self.lora[2](x) + if bias is not None: + x = x + bias + return x + + +class RWKV7GroupNorm(nn.Module): + def __init__( + self, + num_heads: int, + head_dim: int, + value_dim: int, + eps: float, + ) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + self.value_dim = value_dim + self.tp_rank = get_tp_rank() + self.tp_size = get_tp_world_size() + self.local_num_heads = self.num_heads // self.tp_size + self.local_value_dim = self.value_dim // self.tp_size + self.value_start = self.tp_rank * self.local_value_dim + self.value_end = self.value_start + self.local_value_dim + self.eps = self.head_dim * eps + + self.weight = nn.Parameter(torch.ones(self.value_dim)) + self.bias = nn.Parameter(torch.zeros(self.value_dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + weight = self.weight[self.value_start : self.value_end] + bias = self.bias[self.value_start : self.value_end] + x = x.to(torch.float32) + x = F.group_norm( + x.unsqueeze(-1), + num_groups=self.local_num_heads, + weight=weight.to(torch.float32), + bias=bias.to(torch.float32), + eps=self.eps, + ) + return x.squeeze(-1) + + +class RWKV7FeedForward(nn.Module): + def __init__( + self, + config, + layer_idx: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + if config.intermediate_size is None: + hidden_ratio = 4 if config.hidden_ratio is None else config.hidden_ratio + intermediate_size = int(config.hidden_size * hidden_ratio) + intermediate_size = 32 * ((intermediate_size + 31) // 32) + else: + intermediate_size = config.intermediate_size + + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.act_fn = get_activation_fn(config.hidden_act) + self.x_k = nn.Parameter(torch.zeros(self.hidden_size)) + self.key = ColumnParallelLinear( + self.hidden_size, + intermediate_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.key", + ) + self.value = RowParallelLinear( + intermediate_size, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.value", + ) + + def forward( + self, hidden_states: torch.Tensor, cached_state: torch.Tensor | None + ) -> tuple[torch.Tensor, torch.Tensor]: + delta, final_state = token_shift_with_cache(hidden_states, cached_state) + mixed = hidden_states.addcmul(delta, self.x_k) + hidden, _ = self.key(mixed) + hidden = self.act_fn(hidden) + hidden, _ = self.value(hidden) + return hidden, final_state + + def forward_decode_batch( + self, + hidden_states: torch.Tensor, + cached_state: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + delta = cached_state.to(hidden_states.dtype) - hidden_states + mixed = hidden_states.addcmul(delta, self.x_k) + hidden, _ = self.key(mixed) + hidden = self.act_fn(hidden) + hidden, _ = self.value(hidden) + return hidden, hidden_states + + def forward_prefill_batch( + self, + hidden_states: torch.Tensor, + query_start_loc: torch.Tensor, + cached_state: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + delta, final_state = token_shift_with_cache_varlen( + hidden_states, + query_start_loc, + cached_state, + ) + mixed = hidden_states.addcmul(delta, self.x_k) + hidden, _ = self.key(mixed) + hidden = self.act_fn(hidden) + hidden, _ = self.value(hidden) + return hidden, final_state + + +class RWKV7Attention(nn.Module): + def __init__( + self, + config, + layer_idx: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_idx = layer_idx + self.prefix = prefix + self.hidden_size = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = config.head_dim + self.value_dim = config.value_dim[layer_idx] + self.head_v_dim = self.value_dim // self.num_heads + self.tp_rank = get_tp_rank() + self.tp_size = get_tp_world_size() + self.local_num_heads = self.num_heads // self.tp_size + self.local_key_dim = self.hidden_size // self.tp_size + self.local_value_dim = self.value_dim // self.tp_size + self.key_start = self.tp_rank * self.local_key_dim + self.key_end = self.key_start + self.local_key_dim + self.value_start = self.tp_rank * self.local_value_dim + self.value_end = self.value_start + self.local_value_dim + + self.x_r = nn.Parameter(torch.zeros(1, 1, self.hidden_size)) + self.x_w = nn.Parameter(torch.zeros(1, 1, self.hidden_size)) + self.x_k = nn.Parameter(torch.zeros(1, 1, self.hidden_size)) + self.x_v = nn.Parameter(torch.zeros(1, 1, self.hidden_size)) + self.x_a = nn.Parameter(torch.zeros(1, 1, self.hidden_size)) + self.x_g = nn.Parameter(torch.zeros(1, 1, self.hidden_size)) + + self.k_k = nn.Parameter(torch.zeros(self.hidden_size)) + self.k_a = nn.Parameter(torch.zeros(self.hidden_size)) + self.r_k = nn.Parameter(torch.zeros(self.num_heads, self.head_dim)) + + self.r_proj = ColumnParallelLinear( + self.hidden_size, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.r_proj", + ) + self.k_proj = ColumnParallelLinear( + self.hidden_size, + self.hidden_size, + bias=False, + # Quantization error in the key/value projections is written into + # the recurrent state at every token. On larger RWKV7 checkpoints + # this error compounds into non-finite logits. Keep the two state + # update projections in the model dtype; the read-only receptance + # and output projections, FFN, and LM head remain quantizable. + quant_config=None, + prefix=f"{prefix}.k_proj", + ) + self.v_proj = ColumnParallelLinear( + self.hidden_size, + self.value_dim, + bias=False, + quant_config=None, + prefix=f"{prefix}.v_proj", + ) + self.o_proj = RowParallelLinear( + self.value_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.w_lora = RWKV7LoRA( + self.hidden_size, + self.hidden_size, + config.decay_low_rank_dim, + bias=True, + activation="tanh", + quant_config=quant_config, + prefix=f"{prefix}.w_lora", + ) + self.a_lora = RWKV7LoRA( + self.hidden_size, + self.hidden_size, + config.a_low_rank_dim, + bias=True, + activation=None, + quant_config=quant_config, + prefix=f"{prefix}.a_lora", + ) + if self.layer_idx != 0: + self.v_lora = RWKV7LoRA( + self.hidden_size, + self.value_dim, + config.v_low_rank_dim, + bias=True, + activation=None, + quant_config=quant_config, + prefix=f"{prefix}.v_lora", + ) + self.g_lora = RWKV7LoRA( + self.hidden_size, + self.value_dim, + config.gate_low_rank_dim, + bias=False, + activation="sigmoid", + quant_config=quant_config, + prefix=f"{prefix}.g_lora", + ) + self.g_norm = RWKV7GroupNorm( + num_heads=self.num_heads, + head_dim=self.head_dim, + value_dim=self.value_dim, + eps=config.norm_eps, + ) + + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def _project_recurrent_inputs( + self, + hidden_states: torch.Tensor, + delta: torch.Tensor, + v_first: torch.Tensor | None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + x_r = self.x_r.squeeze(0).squeeze(0) + x_w = self.x_w.squeeze(0).squeeze(0) + x_k = self.x_k.squeeze(0).squeeze(0) + x_v = self.x_v.squeeze(0).squeeze(0) + x_a = self.x_a.squeeze(0).squeeze(0) + x_g = self.x_g.squeeze(0).squeeze(0) + + xr = hidden_states.addcmul(delta, x_r) + xw = hidden_states.addcmul(delta, x_w) + xk = hidden_states.addcmul(delta, x_k) + xv = hidden_states.addcmul(delta, x_v) + xa = hidden_states.addcmul(delta, x_a) + xg = hidden_states.addcmul(delta, x_g) + + r, _ = self.r_proj(xr) + w = LOG_DECAY_SCALE * self.w_lora(xw).sigmoid() + k, _ = self.k_proj(xk) + v, _ = self.v_proj(xv) + + if self.layer_idx == 0: + v_first_out = v + else: + if v_first is None: + raise ValueError("RWKV7 layers after layer 0 require `v_first`.") + v = torch.lerp(v, v_first, self.v_lora(xv).sigmoid()) + v_first_out = v_first + + a = self.a_lora(xa).sigmoid() + g = self.g_lora(xg) + + r = r.view(-1, self.local_num_heads, self.head_dim).to(torch.float32) + w = w.view(-1, self.local_num_heads, self.head_dim).to(torch.float32) + k = k.view(-1, self.local_num_heads, self.head_dim).to(torch.float32) + a = a.view(-1, self.local_num_heads, self.head_dim).to(torch.float32) + v = v.view(-1, self.local_num_heads, self.head_v_dim).to(torch.float32) + + local_k_k = self.k_k[self.key_start : self.key_end].view( + 1, self.local_num_heads, self.head_dim + ) + local_k_a = self.k_a[self.key_start : self.key_end].view( + 1, self.local_num_heads, self.head_dim + ) + kk = F.normalize(k * local_k_k.to(torch.float32), dim=-1, p=2.0) + k = k * (1 + (a - 1) * local_k_a.to(torch.float32)) + return r, w, k, v, kk, a, g, v_first_out + + def _finalize_attention_output( + self, + recurrent_output: torch.Tensor, + r: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + hidden_dtype: torch.dtype, + ) -> torch.Tensor: + output = recurrent_output.reshape(-1, self.local_value_dim) + output = self.g_norm(output) + + local_r_k = self.r_k[ + self.tp_rank * self.local_num_heads : (self.tp_rank + 1) + * self.local_num_heads + ].to(torch.float32) + correction = ( + (r * k * local_r_k.unsqueeze(0)).sum(dim=-1, keepdim=True) * v + ).reshape(-1, self.local_value_dim) + output = (output + correction) * g.to(torch.float32) + output = output.to(hidden_dtype) + output, _ = self.o_proj(output) + return output + + def _forward( + self, + hidden_states: torch.Tensor, + cached_shift_state: torch.Tensor | None, + recurrent_state: torch.Tensor | None, + v_first: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + delta, final_shift_state = token_shift_with_cache( + hidden_states, cached_shift_state + ) + r, w, k, v, kk, a, g, v_first_out = self._project_recurrent_inputs( + hidden_states, + delta, + v_first, + ) + + recurrent_output, final_recurrent_state = _rwkv7_recurrent_scan( + r, + w, + k, + v, + kk, + a, + recurrent_state, + ) + + output = self._finalize_attention_output( + recurrent_output, + r, + k, + v, + g, + hidden_states.dtype, + ) + return output, final_shift_state, final_recurrent_state, v_first_out + + def forward( + self, + hidden_states: torch.Tensor, + cached_shift_state: torch.Tensor | None, + recurrent_state: torch.Tensor | None, + v_first: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if not is_forward_context_available() or not hidden_states.is_cuda: + return self._forward( + hidden_states, + cached_shift_state, + recurrent_state, + v_first, + ) + + output = hidden_states.new_empty(hidden_states.shape[0], self.hidden_size) + final_shift_state = hidden_states.new_empty( + self.hidden_size, + dtype=( + cached_shift_state.dtype + if cached_shift_state is not None + else hidden_states.dtype + ), + ) + final_recurrent_state = hidden_states.new_empty( + self.local_num_heads, + self.head_v_dim, + self.head_dim, + dtype=RWKV7_RUNTIME_DTYPE, + ) + v_first_out = hidden_states.new_empty( + hidden_states.shape[0], self.local_value_dim + ) + torch.ops.vllm.rwkv7_attention( + hidden_states, + _custom_op_optional_tensor(cached_shift_state, like=hidden_states), + _custom_op_optional_tensor( + recurrent_state, + like=hidden_states, + dtype=RWKV7_RUNTIME_DTYPE, + ), + _custom_op_optional_tensor(v_first, like=hidden_states), + output, + final_shift_state, + final_recurrent_state, + v_first_out, + self.prefix, + ) + return output, final_shift_state, final_recurrent_state, v_first_out + + def forward_decode_batch( + self, + hidden_states: torch.Tensor, + cached_shift_state: torch.Tensor, + recurrent_state: torch.Tensor, + v_first: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + delta = cached_shift_state.to(hidden_states.dtype) - hidden_states + final_shift_state = hidden_states + r, w, k, v, kk, a, g, v_first_out = self._project_recurrent_inputs( + hidden_states, + delta, + v_first, + ) + recurrent_state = recurrent_state.to(torch.float32) + recurrent_output, final_recurrent_state = _rwkv7_recurrent_step( + r, + w, + k, + v, + kk, + a, + recurrent_state, + ) + + output = self._finalize_attention_output( + recurrent_output, + r, + k, + v, + g, + hidden_states.dtype, + ) + return output, final_shift_state, final_recurrent_state, v_first_out + + def forward_prefill_batch( + self, + hidden_states: torch.Tensor, + query_start_loc: torch.Tensor, + cached_shift_state: torch.Tensor | None, + recurrent_state: torch.Tensor | None, + v_first: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + delta, final_shift_state = token_shift_with_cache_varlen( + hidden_states, + query_start_loc, + cached_shift_state, + ) + r, w, k, v, kk, a, g, v_first_out = self._project_recurrent_inputs( + hidden_states, + delta, + v_first, + ) + + recurrent_output, final_recurrent_state = _rwkv7_recurrent_scan_varlen( + r, + w, + k, + v, + kk, + a, + query_start_loc, + recurrent_state, + ) + output = self._finalize_attention_output( + recurrent_output, + r, + k, + v, + g, + hidden_states.dtype, + ) + return output, final_shift_state, final_recurrent_state, v_first_out + + +class RWKV7Block(nn.Module, MambaBase): + def __init__( + self, + config, + layer_idx: int, + model_config: ModelConfig | None = None, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.model_config = model_config + self.cache_config = cache_config + self.prefix = prefix + self.hidden_size = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = config.head_dim + self.value_dim = config.value_dim[layer_idx] + self.local_value_dim = self.value_dim // get_tp_world_size() + + self.pre_norm = None + if config.norm_first and layer_idx == 0: + self.pre_norm = nn.LayerNorm( + config.hidden_size, + eps=config.norm_eps, + elementwise_affine=True, + bias=config.norm_bias, + ) + self.attn_norm = nn.LayerNorm( + config.hidden_size, + eps=config.norm_eps, + elementwise_affine=True, + bias=config.norm_bias, + ) + self.attn = RWKV7Attention( + config=config, + layer_idx=layer_idx, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + self.ffn_norm = nn.LayerNorm( + config.hidden_size, + eps=config.norm_eps, + elementwise_affine=True, + bias=config.norm_bias, + ) + self.ffn = RWKV7FeedForward( + config=config, + layer_idx=layer_idx, + quant_config=quant_config, + prefix=f"{prefix}.ffn", + ) + + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + self.kv_cache = ( + torch.tensor([]), + torch.tensor([]), + torch.tensor([]), + ) + + @property + def mamba_type(self) -> MambaAttentionBackendEnum: + return MambaAttentionBackendEnum.LINEAR + + def get_state_dtype(self) -> tuple[torch.dtype, ...]: + return ( + RWKV7_RUNTIME_DTYPE, + RWKV7_RUNTIME_DTYPE, + RWKV7_RUNTIME_DTYPE, + ) + + def get_state_shape(self) -> tuple[tuple[int, ...], ...]: + return ( + (self.hidden_size,), + ( + self.num_heads // get_tp_world_size(), + self.value_dim // self.num_heads, + self.head_dim, + ), + (self.hidden_size,), + ) + + def _run_sequence( + self, + hidden_states: torch.Tensor, + v_first: torch.Tensor | None, + attn_shift_state: torch.Tensor | None, + recurrent_state: torch.Tensor | None, + ffn_shift_state: torch.Tensor | None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + residual = hidden_states + if self.pre_norm is not None: + residual = self.pre_norm(residual) + + attn_input = self.attn_norm(residual) + attn_out, attn_shift_state, recurrent_state, v_first_out = self.attn( + attn_input, + attn_shift_state, + recurrent_state, + v_first, + ) + hidden_states = residual + attn_out + + ffn_input = self.ffn_norm(hidden_states) + ffn_out, ffn_shift_state = self.ffn(ffn_input, ffn_shift_state) + hidden_states = hidden_states + ffn_out + return ( + hidden_states, + v_first_out, + attn_shift_state, + recurrent_state, + ffn_shift_state, + ) + + def _get_kv_state( + self, slot_id: int, use_initial_state: bool + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + if not use_initial_state: + return None, None, None + return ( + self.kv_cache[0][slot_id], + self.kv_cache[1][slot_id], + self.kv_cache[2][slot_id], + ) + + def _get_kv_states( + self, slot_ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return ( + self.kv_cache[0].index_select(0, slot_ids), + self.kv_cache[1].index_select(0, slot_ids), + self.kv_cache[2].index_select(0, slot_ids), + ) + + def _get_prefill_kv_states( + self, + slot_ids: torch.Tensor, + has_initial_state: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + attn_shift_state, recurrent_state, ffn_shift_state = self._get_kv_states( + slot_ids + ) + has_initial_state = has_initial_state.to(attn_shift_state.device) + attn_shift_state = torch.where( + has_initial_state[:, None], + attn_shift_state, + torch.zeros_like(attn_shift_state), + ) + recurrent_state = torch.where( + has_initial_state[:, None, None, None], + recurrent_state, + torch.zeros_like(recurrent_state), + ) + ffn_shift_state = torch.where( + has_initial_state[:, None], + ffn_shift_state, + torch.zeros_like(ffn_shift_state), + ) + return attn_shift_state, recurrent_state, ffn_shift_state + + @torch.compiler.disable + def _store_kv_state( + self, + slot_id: int, + attn_shift_state: torch.Tensor, + recurrent_state: torch.Tensor, + ffn_shift_state: torch.Tensor, + ) -> None: + self.kv_cache[0][slot_id].copy_( + attn_shift_state.to(self.kv_cache[0][slot_id].dtype) + ) + self.kv_cache[1][slot_id].copy_( + recurrent_state.to(self.kv_cache[1][slot_id].dtype) + ) + self.kv_cache[2][slot_id].copy_( + ffn_shift_state.to(self.kv_cache[2][slot_id].dtype) + ) + + @torch.compiler.disable + def _store_kv_states( + self, + slot_ids: torch.Tensor, + attn_shift_state: torch.Tensor, + recurrent_state: torch.Tensor, + ffn_shift_state: torch.Tensor, + ) -> None: + self.kv_cache[0].index_copy_( + 0, slot_ids, attn_shift_state.to(self.kv_cache[0].dtype) + ) + self.kv_cache[1].index_copy_( + 0, slot_ids, recurrent_state.to(self.kv_cache[1].dtype) + ) + self.kv_cache[2].index_copy_( + 0, slot_ids, ffn_shift_state.to(self.kv_cache[2].dtype) + ) + + def _run_decode_batch( + self, + hidden_states: torch.Tensor, + v_first: torch.Tensor | None, + attn_shift_state: torch.Tensor, + recurrent_state: torch.Tensor, + ffn_shift_state: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + residual = hidden_states + if self.pre_norm is not None: + residual = self.pre_norm(residual) + + attn_input = self.attn_norm(residual) + attn_out, attn_shift_state, recurrent_state, v_first_out = ( + self.attn.forward_decode_batch( + attn_input, + attn_shift_state, + recurrent_state, + v_first, + ) + ) + hidden_states = residual + attn_out + + ffn_input = self.ffn_norm(hidden_states) + ffn_out, ffn_shift_state = self.ffn.forward_decode_batch( + ffn_input, ffn_shift_state + ) + hidden_states = hidden_states + ffn_out + return ( + hidden_states, + v_first_out, + attn_shift_state, + recurrent_state, + ffn_shift_state, + ) + + def _run_prefill_batch( + self, + hidden_states: torch.Tensor, + v_first: torch.Tensor | None, + query_start_loc: torch.Tensor, + attn_shift_state: torch.Tensor | None, + recurrent_state: torch.Tensor | None, + ffn_shift_state: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + residual = hidden_states + if self.pre_norm is not None: + residual = self.pre_norm(residual) + + attn_input = self.attn_norm(residual) + attn_out, attn_shift_state, recurrent_state, v_first_out = ( + self.attn.forward_prefill_batch( + attn_input, + query_start_loc, + attn_shift_state, + recurrent_state, + v_first, + ) + ) + hidden_states = residual + attn_out + + ffn_input = self.ffn_norm(hidden_states) + ffn_out, ffn_shift_state = self.ffn.forward_prefill_batch( + ffn_input, + query_start_loc, + ffn_shift_state, + ) + hidden_states = hidden_states + ffn_out + return ( + hidden_states, + v_first_out, + attn_shift_state, + recurrent_state, + ffn_shift_state, + ) + + def _forward_runtime( + self, + hidden_states: torch.Tensor, + v_first: torch.Tensor | None, + output: torch.Tensor, + v_first_out: torch.Tensor, + attn_metadata: LinearAttentionMetadata | None = None, + ) -> None: + if attn_metadata is None and is_forward_context_available(): + forward_context = get_forward_context() + runtime_attn_metadata = forward_context.attn_metadata + if runtime_attn_metadata is not None: + assert isinstance(runtime_attn_metadata, dict) + maybe_metadata = runtime_attn_metadata.get(self.prefix) + assert maybe_metadata is None or isinstance( + maybe_metadata, LinearAttentionMetadata + ) + attn_metadata = maybe_metadata + + if attn_metadata is None: + out, vf_out, _, _, _ = self._run_sequence( + hidden_states, v_first, None, None, None + ) + output[: out.shape[0]] = out + v_first_out[: vf_out.shape[0]] = vf_out + return + + num_actual_tokens = ( + attn_metadata.num_decode_tokens + attn_metadata.num_prefill_tokens + ) + hidden_states = hidden_states[:num_actual_tokens] + if v_first is not None: + v_first = v_first[:num_actual_tokens] + + output_slice = output[:num_actual_tokens] + v_first_slice = v_first_out[:num_actual_tokens] + state_indices = attn_metadata.state_indices_tensor + + if attn_metadata.num_decode_tokens > 0: + # Full CUDA graphs pad the request batch with zero-length decode + # requests. ``num_decodes`` includes those padded requests, while + # ``num_decode_tokens`` only counts real one-token decodes. Avoid + # indexing the state cache with the padded ``PAD_SLOT_ID`` entries. + decode_slot_ids = state_indices[: attn_metadata.num_decode_tokens].to( + dtype=torch.long + ) + states = self._get_kv_states(decode_slot_ids) + out, vf_out, attn_shift, recurrent, ffn_shift = self._run_decode_batch( + hidden_states[: attn_metadata.num_decode_tokens], + None if v_first is None else v_first[: attn_metadata.num_decode_tokens], + *states, + ) + output_slice[: attn_metadata.num_decode_tokens] = out + v_first_slice[: attn_metadata.num_decode_tokens] = vf_out + self._store_kv_states( + decode_slot_ids, + attn_shift, + recurrent, + ffn_shift, + ) + + prefill_req_offset = attn_metadata.num_decodes + prefill_token_offset = attn_metadata.num_decode_tokens + if attn_metadata.num_prefills > 0: + prefill_req_end = prefill_req_offset + attn_metadata.num_prefills + prefill_slot_ids = state_indices[prefill_req_offset:prefill_req_end].to( + dtype=torch.long + ) + prefill_query_start_loc = ( + attn_metadata.query_start_loc[prefill_req_offset : prefill_req_end + 1] + - prefill_token_offset + ) + query_lens = prefill_query_start_loc[1:] - prefill_query_start_loc[:-1] + prefill_seq_lens = attn_metadata.seq_lens[ + prefill_req_offset:prefill_req_end + ] + has_initial_state = prefill_seq_lens > query_lens + states = self._get_prefill_kv_states( + prefill_slot_ids, + has_initial_state, + ) + out, vf_out, attn_shift, recurrent, ffn_shift = self._run_prefill_batch( + hidden_states[prefill_token_offset:num_actual_tokens], + ( + None + if v_first is None + else v_first[prefill_token_offset:num_actual_tokens] + ), + prefill_query_start_loc, + *states, + ) + output_slice[prefill_token_offset:num_actual_tokens] = out + v_first_slice[prefill_token_offset:num_actual_tokens] = vf_out + self._store_kv_states( + prefill_slot_ids, + attn_shift, + recurrent, + ffn_shift, + ) + + def forward( + self, + hidden_states: torch.Tensor, + v_first: torch.Tensor | None, + attn_metadata: LinearAttentionMetadata | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + output = torch.empty_like(hidden_states) + v_first_out = torch.empty( + (hidden_states.shape[0], self.local_value_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + if is_forward_context_available() and hidden_states.is_cuda: + torch.ops.vllm.rwkv7_block_forward( + hidden_states, + _custom_op_optional_tensor(v_first, like=hidden_states), + output, + v_first_out, + self.prefix, + ) + else: + self._forward_runtime( + hidden_states, + v_first, + output, + v_first_out, + attn_metadata=attn_metadata, + ) + return output, v_first_out + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": 0, + "intermediate_tensors": 0, + "inputs_embeds": 0, + } +) +class RWKV7Model(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_config + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.config = config + self.runtime_dtype = model_config.dtype + + if config.attn is not None: + raise NotImplementedError( + "Hybrid RWKV7 checkpoints with transformer attention are " + "not supported yet." + ) + + value_dims = config.value_dim + if len(set(value_dims)) != 1: + raise NotImplementedError( + "RWKV7 with per-layer `value_dim` variation is not supported yet." + ) + + self.local_value_dim = value_dims[0] // get_tp_world_size() + self.vocab_size = config.vocab_size + self.embed_tokens = ( + VocabParallelEmbedding(config.vocab_size, config.hidden_size) + if get_pp_group().is_first_rank + else PPMissingLayer() + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: RWKV7Block( + config=config, + layer_idx=int(prefix.split(".")[-1]), + model_config=model_config, + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + ), + prefix=f"{prefix}.layers", + ) + + self.norm = ( + nn.LayerNorm( + config.hidden_size, + eps=config.norm_eps, + elementwise_affine=True, + bias=config.norm_bias, + ) + if get_pp_group().is_last_rank + else PPMissingLayer() + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def _pp_runtime_dtype(self) -> torch.dtype: + # RWKV7 keeps inter-stage activations in the model dtype so PP dummy runs + # and stage-to-stage transfers match the numerics used inside blocks. + return self.runtime_dtype + + def make_empty_intermediate_tensors( + self, batch_size: int, dtype: torch.dtype, device: torch.device + ) -> IntermediateTensors: + del dtype + runtime_dtype = self._pp_runtime_dtype() + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.config.hidden_size), + dtype=runtime_dtype, + device=device, + ), + "v_first": torch.zeros( + (batch_size, self.local_value_dim), + dtype=runtime_dtype, + device=device, + ), + } + ) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors: + del positions + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + if attn_metadata is not None: + assert isinstance(attn_metadata, dict) + + if get_pp_group().is_first_rank: + hidden_states = ( + inputs_embeds + if inputs_embeds is not None + else self.embed_input_ids(input_ids) + ) + v_first = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + v_first = intermediate_tensors["v_first"] + runtime_dtype = self._pp_runtime_dtype() + if hidden_states.dtype != runtime_dtype: + hidden_states = hidden_states.to(runtime_dtype) + if v_first.dtype != runtime_dtype: + v_first = v_first.to(runtime_dtype) + + for layer in islice(self.layers, self.start_layer, self.end_layer): + layer_metadata = ( + None if attn_metadata is None else attn_metadata.get(layer.prefix) + ) + assert layer_metadata is None or isinstance( + layer_metadata, LinearAttentionMetadata + ) + hidden_states, v_first = layer( + hidden_states=hidden_states, + v_first=v_first, + attn_metadata=layer_metadata, + ) + + if not get_pp_group().is_last_rank: + assert v_first is not None + runtime_dtype = self._pp_runtime_dtype() + if hidden_states.dtype != runtime_dtype: + hidden_states = hidden_states.to(runtime_dtype) + if v_first.dtype != runtime_dtype: + v_first = v_first.to(runtime_dtype) + return IntermediateTensors( + {"hidden_states": hidden_states, "v_first": v_first} + ) + + hidden_states = self.norm(hidden_states) + return hidden_states + + +class RWKV7ForCausalLM( + nn.Module, + HasInnerState, + IsAttentionFree, + SupportsPP, +): + # RWKV7 keeps the checkpoint's linear projections separate, so there are + # no packed-module rewrites to describe. The attribute is still required + # by weight-rewriting loaders such as BitsAndBytes; an explicit empty map + # declares that checkpoint and vLLM module names match one-to-one. + packed_modules_mapping: dict[str, list[str]] = {} + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_config + self.config = config + self.model = RWKV7Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + + if get_pp_group().is_last_rank: + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + else: + self.lm_head = PPMissingLayer() + + model_dtype = vllm_config.model_config.dtype + self.model.to(model_dtype) + if get_pp_group().is_last_rank: + self.lm_head.to(model_dtype) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def make_empty_intermediate_tensors( + self, batch_size: int, dtype: torch.dtype, device: torch.device + ) -> IntermediateTensors: + return self.model.make_empty_intermediate_tensors(batch_size, dtype, device) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + weight = getattr(self.lm_head, "weight", None) + if weight is not None and weight.is_floating_point(): + hidden_states = hidden_states.to(weight.dtype) + return self.logits_processor(self.lm_head, hidden_states) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[torch.dtype, ...]: + return ( + RWKV7_RUNTIME_DTYPE, + RWKV7_RUNTIME_DTYPE, + RWKV7_RUNTIME_DTYPE, + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[tuple[int, ...], ...]: + config = vllm_config.model_config.hf_config + if len(set(config.value_dim)) != 1: + raise NotImplementedError( + "RWKV7 with per-layer `value_dim` variation is not supported yet." + ) + return ( + (config.hidden_size,), + ( + config.num_heads // vllm_config.parallel_config.tensor_parallel_size, + config.head_dim, + config.value_dim[0] // config.num_heads, + ), + (config.hidden_size,), + ) + + @classmethod + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]: + conv, temporal = MambaStateCopyFuncCalculator.mamba1_state_copy_func() + return (conv, temporal, conv) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + def iter_weights(): + for name, tensor in weights: + if name == "model.embeddings.weight": + yield "model.embed_tokens.weight", tensor + else: + yield name, tensor + + loader = AutoWeightsLoader(self) + return loader.load_weights(iter_weights()) diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 6f8c25769ef2..814430f54382 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -125,6 +125,7 @@ def __getitem__(self, key): qwen3_next="Qwen3NextConfig", qwen3_5="Qwen3_5Config", qwen3_5_moe="Qwen3_5MoeConfig", + rwkv7="RWKV7Config", laguna="LagunaConfig", lfm2_moe="Lfm2MoeConfig", **{"unlimited-ocr": "UnlimitedOCRConfig"}, diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 4bb7674ddbb4..64067200a9cc 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -98,6 +98,7 @@ "InklingAudioConfig": "vllm.models.inkling.configs", "InklingVisionConfig": "vllm.models.inkling.configs", "InklingMMConfig": "vllm.models.inkling.configs", + "RWKV7Config": "vllm.transformers_utils.configs.rwkv7", # Special case: DeepseekV3Config is from HuggingFace Transformers "DeepseekV3Config": "transformers", } @@ -182,6 +183,7 @@ "InklingAudioConfig", "InklingVisionConfig", "InklingMMConfig", + "RWKV7Config", ] diff --git a/vllm/transformers_utils/configs/rwkv7.py b/vllm/transformers_utils/configs/rwkv7.py new file mode 100644 index 000000000000..7b9e1f6f9f34 --- /dev/null +++ b/vllm/transformers_utils/configs/rwkv7.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from transformers.configuration_utils import PretrainedConfig + + +class RWKV7Config(PretrainedConfig): + model_type = "rwkv7" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + num_hidden_layers: int = 24, + head_dim: int | None = 64, + num_heads: int | None = None, + decay_low_rank_dim: int = 64, + gate_low_rank_dim: int = 128, + a_low_rank_dim: int = 64, + v_low_rank_dim: int = 16, + hidden_act: str = "sqrelu", + max_position_embeddings: int = 2048, + norm_first: bool = True, + norm_bias: bool = True, + norm_eps: float = 1e-5, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = True, + vocab_size: int = 32000, + value_dim: int | list[int] | None = None, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.norm_first = norm_first + self.num_hidden_layers = num_hidden_layers + + if head_dim is None and num_heads is not None: + head_dim = int(hidden_size // num_heads) + elif head_dim is not None and num_heads is None: + num_heads = int(hidden_size // head_dim) + elif head_dim is None and num_heads is None: + raise ValueError("Either `head_dim` or `num_heads` must be specified.") + + if value_dim is None: + value_dim = [hidden_size] * num_hidden_layers + elif isinstance(value_dim, int): + if value_dim < hidden_size or value_dim % hidden_size != 0: + raise ValueError( + "`value_dim` must be >= hidden_size and divisible by hidden_size." + ) + value_dim = [value_dim] * num_hidden_layers + else: + if len(value_dim) != num_hidden_layers: + raise ValueError( + "`value_dim` must have the same length as num_hidden_layers." + ) + for dim in value_dim: + if dim < hidden_size or dim % hidden_size != 0: + raise ValueError( + "`value_dim` must be >= hidden_size and divisible " + "by hidden_size." + ) + + self.head_dim = head_dim + self.num_heads = num_heads + self.value_dim = value_dim + self.decay_low_rank_dim = decay_low_rank_dim + self.gate_low_rank_dim = gate_low_rank_dim + self.a_low_rank_dim = a_low_rank_dim + self.v_low_rank_dim = v_low_rank_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.norm_bias = norm_bias + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + self.fuse_norm = fuse_norm + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` " + "cannot both be enabled." + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("`attn` must be a dictionary.") + if "layers" not in attn: + raise ValueError("`attn.layers` must be provided for hybrid RWKV7.") + if "num_heads" not in attn: + raise ValueError("`attn.num_heads` must be provided for hybrid RWKV7.") + attn["num_kv_heads"] = attn.get("num_kv_heads", attn["num_heads"]) + attn["qkv_bias"] = attn.get("qkv_bias", False) + attn["window_size"] = attn.get("window_size", None) + attn["rope_theta"] = attn.get("rope_theta", 10000.0) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + )