diff --git a/README.md b/README.md index 05397d22..39f8b2f0 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,7 @@ Read the method overview and integration rules in | Qwen3 | ✅ | | Qwen3MoE | ✅ | | Qwen3.5 / Qwen3.6 | ✅ | +| Qwen3.5 / Qwen3.6 MoE | ✅ | | Llama 3 / 3.1 | ✅ | | MiniMax M2.7 | ✅ | @@ -158,8 +159,6 @@ the smaller CUDA-specific extra: pip install -e ".[prefix-offload]" ``` -Sparse-vLLM supports Qwen3.5/Qwen3.6 checkpoints in unquantized BF16 and -block-scaled FP8 formats. The Qwen3.5/Qwen3.6 prefill causal Conv1D and decode Conv1D/GDN packing paths use repository-local Triton kernels; `sglang-kernel` and a local CUDA extension diff --git a/README_zh.md b/README_zh.md index 306a5ffc..021af344 100644 --- a/README_zh.md +++ b/README_zh.md @@ -58,6 +58,7 @@ Sparse-vLLM 支持物理淘汰、逻辑掩码、查询感知选择和混合 KV | Qwen3 | ✅ | | Qwen3MoE | ✅ | | Qwen3.5 / Qwen3.6 | ✅ | +| Qwen3.5 / Qwen3.6 MoE | ✅ | | Llama 3 / 3.1 | ✅ | | MiniMax M2.7 | ✅ | @@ -127,7 +128,7 @@ uv pip install flashinfer-cubin --index-url https://flashinfer.ai/whl 以上显式索引分别用于安装 CUDA 13.0 版本的 PyTorch 和 FlashInfer JIT 缓存。 -Qwen3.5/Qwen3.6 混合注意力推理还需要安装 CUDA 专用的可选依赖： +Qwen3.5/Qwen3.6 混合注意力推理还需要安装可选 Python 依赖： ```bash # uv @@ -137,6 +138,17 @@ uv pip install -e ".[qwen35]" pip install -e ".[qwen35]" ``` +对于不使用 Qwen3.5/Qwen3.6 的 Vanilla、OmniKV 或 QuEST 前缀缓存卸载， +可以安装更精简的 CUDA 专用可选依赖： + +```bash +pip install -e ".[prefix-offload]" +``` + +Qwen3.5/Qwen3.6 的 prefill causal Conv1D 和 decode Conv1D/GDN packing +路径使用仓库内置的 Triton kernel；无需安装 `sglang-kernel`，也无需编译 +本地 CUDA 扩展。 + 完整依赖列表和最小 `LLM(...)` 示例请参阅[快速开始](docs/zh/getting_started/README.md)。 ## 基准测试 diff --git a/benchmark/microbench.py b/benchmark/microbench.py index ffa5bf76..765aeaf7 100644 --- a/benchmark/microbench.py +++ b/benchmark/microbench.py @@ -644,7 +644,6 @@ def add_wave(max_new_requests: int): ): ttft = perf_counter() - t_start elif num_tokens < 0: - # print(f'one decode step ... {perf_counter() - last_time}') decode_started = True decode_steps_since_last_wave += 1 decode_times.append(step_dt) @@ -680,12 +679,7 @@ def add_wave(max_new_requests: int): f"{decode_warmup_steps_after_full} warmup steps." ) - print(f'@@@ {decode_tokens=}') - torch.cuda.synchronize() - t_end = perf_counter() - - duration = t_end - t_start peak_mem = get_peak_memory() graph_status = _decode_cuda_graph_status(llm) prefix_cache_stats_after = _cache_stats(llm) @@ -722,7 +716,6 @@ def add_wave(max_new_requests: int): prefill_s = sum(prefill_times) decode_s = sum(decode_times) - print(f'[debug] {prefill_tokens=} {prefill_s=} {ttft=} {decode_tokens=} {decode_s=} {has_queued=}') prefill_tp = prefill_tokens / prefill_s if prefill_s > 0 else 0 used_full_admission_window = bool(decode_times_after_full) decode_s_effective = sum(decode_times_after_full) if used_full_admission_window else decode_s diff --git a/benchmark/sparsevllm_regression/manifest.json b/benchmark/sparsevllm_regression/manifest.json index a9f1483f..803b8f64 100644 --- a/benchmark/sparsevllm_regression/manifest.json +++ b/benchmark/sparsevllm_regression/manifest.json @@ -75,7 +75,7 @@ "h2o": { "sparse_method": "h2o", "requires_compressor": false, - "supported_model_families": ["qwen2", "qwen3", "qwen3_moe", "qwen3_5", "llama", "minimax_m2"], + "supported_model_families": ["qwen2", "qwen3", "qwen3_moe", "qwen3_5", "qwen3_5_moe", "llama", "minimax_m2"], "supported_tensor_parallel_sizes": [1, 2], "config": { "sparse_method": "h2o", diff --git a/benchmark/vllm_microbench.py b/benchmark/vllm_microbench.py new file mode 100644 index 00000000..fa755e39 --- /dev/null +++ b/benchmark/vllm_microbench.py @@ -0,0 +1,282 @@ +"""Reproducible vLLM latency baseline for Sparse-vLLM comparisons. + +Run this script with an isolated vLLM environment. It intentionally imports +vLLM inside ``main`` so the Sparse-vLLM project environment does not need vLLM. +""" + +from __future__ import annotations + +import argparse +import json +import os +import shlex +import statistics +import sys +import traceback +from datetime import datetime +from importlib.metadata import version +from pathlib import Path +from time import perf_counter +from typing import Any + + +def _write_json(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + for row in rows: + json.dump(row, handle, ensure_ascii=False, sort_keys=True) + handle.write("\n") + + +def _parse_positive_ints(value: str) -> list[int]: + values = [int(part.strip()) for part in value.split(",") if part.strip()] + if not values or any(item <= 0 for item in values): + raise argparse.ArgumentTypeError("expected comma-separated positive integers") + if len(values) != len(set(values)): + raise argparse.ArgumentTypeError("batch sizes must be unique") + return values + + +def _env_snapshot() -> dict[str, str]: + keys = ( + "CUDA_VISIBLE_DEVICES", + "VLLM_ALL2ALL_BACKEND", + "VLLM_USE_V1", + "NCCL_DEBUG", + ) + return {key: os.environ[key] for key in keys if key in os.environ} + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Run a fixed-token vLLM baseline compatible with microbench.py." + ) + parser.add_argument("--model-path", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--input-len", type=int, default=1024) + parser.add_argument("--output-len", type=int, default=128) + parser.add_argument("--batch-sizes", type=_parse_positive_ints, default=[1, 2, 4]) + parser.add_argument("--num-warmups", type=int, default=2) + parser.add_argument("--num-iters", type=int, default=5) + parser.add_argument("--tensor-parallel-size", type=int, default=2) + parser.add_argument("--enable-expert-parallel", action="store_true") + parser.add_argument("--gpu-memory-utilization", type=float, default=0.70) + parser.add_argument("--max-model-len", type=int, default=1252) + parser.add_argument("--max-num-batched-tokens", type=int, default=4096) + parser.add_argument("--prompt-token-id", type=int, default=100) + return parser + + +def _validate_args(args: argparse.Namespace) -> None: + positive_names = ( + "input_len", + "output_len", + "num_warmups", + "num_iters", + "tensor_parallel_size", + "max_model_len", + "max_num_batched_tokens", + ) + for name in positive_names: + if int(getattr(args, name)) <= 0: + raise ValueError(f"{name} must be positive") + if args.input_len + args.output_len > args.max_model_len: + raise ValueError( + "max_model_len must cover input_len + output_len: " + f"{args.max_model_len} < {args.input_len + args.output_len}" + ) + if not 0.0 < args.gpu_memory_utilization <= 1.0: + raise ValueError("gpu_memory_utilization must be in (0, 1]") + + +def main() -> int: + args = _build_parser().parse_args() + _validate_args(args) + output_dir = Path(args.output_dir).expanduser().resolve() + if output_dir.exists() and any(output_dir.iterdir()): + raise FileExistsError(f"output directory is not empty: {output_dir}") + output_dir.mkdir(parents=True, exist_ok=True) + + batch_sizes = list(args.batch_sizes) + engine_config = { + "model": str(Path(args.model_path).expanduser().resolve()), + "tensor_parallel_size": int(args.tensor_parallel_size), + "enable_expert_parallel": bool(args.enable_expert_parallel), + "gpu_memory_utilization": float(args.gpu_memory_utilization), + "max_model_len": int(args.max_model_len), + "max_num_seqs": max(batch_sizes), + "max_num_batched_tokens": int(args.max_num_batched_tokens), + "enable_prefix_caching": False, + "language_model_only": True, + "seed": 0, + "enforce_eager": False, + "disable_log_stats": True, + "compilation_config": { + "cudagraph_capture_sizes": batch_sizes, + "max_cudagraph_capture_size": max(batch_sizes), + }, + } + run_info = { + "benchmark": "vllm_microbench", + "created_at": datetime.now().isoformat(timespec="seconds"), + "command": shlex.join(sys.argv), + "engine_config": engine_config, + "input_len": int(args.input_len), + "output_len": int(args.output_len), + "batch_sizes": batch_sizes, + "num_warmups": int(args.num_warmups), + "num_iters": int(args.num_iters), + "prompt_token_id": int(args.prompt_token_id), + "sampling": { + "temperature": 0.0, + "top_p": 1.0, + "ignore_eos": True, + }, + "env": _env_snapshot(), + } + _write_json(output_dir / "run_info.json", run_info) + + performance_rows: list[dict[str, Any]] = [] + per_sample_rows: list[dict[str, Any]] = [] + raw_output_rows: list[dict[str, Any]] = [] + llm = None + try: + import torch + import transformers + import vllm + from vllm import LLM, SamplingParams + + run_info["versions"] = { + "vllm": vllm.__version__, + "torch": torch.__version__, + "transformers": transformers.__version__, + "flashinfer_python": version("flashinfer-python"), + } + _write_json(output_dir / "run_info.json", run_info) + + llm = LLM(**engine_config) + sampling_params = SamplingParams( + temperature=0.0, + top_p=1.0, + ignore_eos=True, + max_tokens=int(args.output_len), + detokenize=False, + ) + + for batch_size in batch_sizes: + prompts = [ + {"prompt_token_ids": [int(args.prompt_token_id)] * args.input_len} + for _ in range(batch_size) + ] + for _ in range(args.num_warmups): + warmup_outputs = llm.generate( + prompts, + sampling_params=sampling_params, + use_tqdm=False, + ) + if len(warmup_outputs) != batch_size: + raise RuntimeError( + f"warmup returned {len(warmup_outputs)} requests, " + f"expected {batch_size}" + ) + + latencies: list[float] = [] + for iteration in range(args.num_iters): + started = perf_counter() + outputs = llm.generate( + prompts, + sampling_params=sampling_params, + use_tqdm=False, + ) + latency = perf_counter() - started + latencies.append(latency) + if len(outputs) != batch_size: + raise RuntimeError( + f"iteration {iteration} returned {len(outputs)} requests, " + f"expected {batch_size}" + ) + + for sample_index, output in enumerate(outputs): + if len(output.outputs) != 1: + raise RuntimeError( + f"iteration {iteration} sample {sample_index} returned " + f"{len(output.outputs)} sequences, expected 1" + ) + token_ids = list(output.outputs[0].token_ids) + status = ( + "success" + if len(token_ids) == args.output_len + else "model_failed" + ) + sample_row = { + "batch_size": batch_size, + "iteration": iteration, + "sample_index": sample_index, + "status": status, + "input_tokens": int(args.input_len), + "output_tokens": len(token_ids), + } + per_sample_rows.append(sample_row) + raw_output_rows.append({**sample_row, "token_ids": token_ids}) + if status != "success": + raise RuntimeError( + f"iteration {iteration} sample {sample_index} produced " + f"{len(token_ids)} tokens, expected {args.output_len}" + ) + + mean_latency = statistics.fmean(latencies) + performance_rows.append( + { + "batch_size": batch_size, + "status": "success", + "latencies_s": latencies, + "e2e_latency_s_mean": mean_latency, + "e2e_latency_s_median": statistics.median(latencies), + "input_tok_s": batch_size * args.input_len / mean_latency, + "output_tok_s": batch_size * args.output_len / mean_latency, + "total_tok_s": ( + batch_size * (args.input_len + args.output_len) / mean_latency + ), + } + ) + + aggregate = { + "benchmark": "vllm_microbench", + "status": "success", + "num_cases": len(performance_rows), + "records": performance_rows, + } + except Exception as error: + aggregate = { + "benchmark": "vllm_microbench", + "status": "model_failed", + "error": repr(error), + "traceback": traceback.format_exc(), + "records": performance_rows, + } + _write_jsonl(output_dir / "raw_outputs.jsonl", raw_output_rows) + _write_jsonl(output_dir / "per_sample_results.jsonl", per_sample_rows) + _write_jsonl(output_dir / "performance.jsonl", performance_rows) + _write_json(output_dir / "aggregate_metrics.json", aggregate) + raise + finally: + if llm is not None: + del llm + + _write_jsonl(output_dir / "raw_outputs.jsonl", raw_output_rows) + _write_jsonl(output_dir / "per_sample_results.jsonl", per_sample_rows) + _write_jsonl(output_dir / "performance.jsonl", performance_rows) + _write_json(output_dir / "aggregate_metrics.json", aggregate) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/configs/debug/deepseek_v4_flash_tiny_random.json b/configs/debug/deepseek_v4_flash_tiny_random.json new file mode 100644 index 00000000..144965a0 --- /dev/null +++ b/configs/debug/deepseek_v4_flash_tiny_random.json @@ -0,0 +1,18 @@ +{ + "head_dim": 64, + "hidden_size": 128, + "index_head_dim": 64, + "index_n_heads": 4, + "index_topk": 8, + "intermediate_size": 64, + "n_routed_experts": 8, + "num_attention_heads": 4, + "num_experts_per_tok": 2, + "num_hidden_layers": 4, + "num_local_experts": 8, + "o_groups": 2, + "o_lora_rank": 64, + "q_lora_rank": 64, + "sliding_window": 8, + "vocab_size": 128 +} diff --git a/docs/en/features/supported-models.md b/docs/en/features/supported-models.md index 7bb3942a..8ff710a8 100644 --- a/docs/en/features/supported-models.md +++ b/docs/en/features/supported-models.md @@ -16,6 +16,7 @@ parallel size must use that value. | Qwen3 Dense | `qwen3` | BF16 / FP16 / block FP8 | ✅ (FP8: 1/2/4/8) | 1 only | 1 only | | Qwen3MoE | `qwen3_moe` | BF16 / FP16 / block FP8 | ✅ (TP > 1: BF16 model dtype only) | 1 only | ✅ | | Qwen3.5 / Qwen3.6 | `qwen3_5` | BF16 / block FP8 | ✅ | 1 only | 1 only | +| Qwen3.6 MoE | `qwen3_5_moe` | BF16 / block FP8 | ✅ | 1 only | ✅ | | Llama 3 / 3.1 | `llama` | BF16 / FP16 | ✅ | 1 only | 1 only | | MiniMax M2.7 | `minimax_m2` | block FP8 with BF16 non-quantized weights | ✅ | 1 only | ✅ | @@ -32,8 +33,9 @@ model dtype; FP16 Qwen3MoE checkpoints are limited to `TP=1`. When `TP=1`, the existing EP layout uses world size `E`. Block FP8 support requires E4M3 weights, dynamic activation quantization, and -a `128 x 128` weight block size. Qwen3.5 and Qwen3.6 configurations are -normalized internally to `model_type=qwen3_5`. +a `128 x 128` weight block size. Qwen3.5/Qwen3.6 dense configurations are +normalized internally to `model_type=qwen3_5`; Qwen3.6 MoE uses +`model_type=qwen3_5_moe`. ## Sparse Method Support @@ -43,6 +45,7 @@ normalized internally to `model_type=qwen3_5`. | Qwen3 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | Compressor required² | | Qwen3MoE | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Qwen3.5 / Qwen3.6 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | Matched checkpoint³ | +| Qwen3.6 MoE | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Llama 3 / 3.1 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | Selected checkpoint¹ | Compressor required² | | MiniMax M2.7 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | diff --git a/docs/zh/benchmarking/README.md b/docs/zh/benchmarking/README.md index b8d9202a..1b1c3182 100644 --- a/docs/zh/benchmarking/README.md +++ b/docs/zh/benchmarking/README.md @@ -20,6 +20,7 @@ | RULER-VT | `benchmark/ruler_vt/pred.py` | 使用原生 benchmark adapter、自包含的 RULER variable-tracking generator/evaluator。 | | NIAH | `benchmark/niah/test_niah.py` | 原生 Sparse-vLLM needle-in-a-haystack 长上下文 runner。 | | Regression harness | [`sparsevllm-regression-tests.md`](sparsevllm-regression-tests.md) | 固定的 quality/performance/stress 检查。 | +| DeepSeek V4 Flash | [`deepseek-v4-flash.md`](deepseek-v4-flash.md) | DPA+EP、CUDA Graph、正确性与性能验收记录。 | ## 吞吐量 Benchmark diff --git a/docs/zh/benchmarking/deepseek-v4-flash.md b/docs/zh/benchmarking/deepseek-v4-flash.md new file mode 100644 index 00000000..5538184c --- /dev/null +++ b/docs/zh/benchmarking/deepseek-v4-flash.md @@ -0,0 +1,110 @@ +# DeepSeek V4 Flash 验证与性能数据 + +本文记录 DeepSeek V4 Flash 首版适配的可复现实验结果。所有 GPU 结果只使用物理 GPU 4–7 中执行前确认空闲的设备。 + +## 正确性 + +| 检查 | 结果 | +| --- | --- | +| Sliding attention，13-token prefill | 对 Transformers reference 最大绝对误差 `1.19e-7` | +| CSA，13-token prefill | 对 Transformers reference 最大绝对误差 `1.19e-7` | +| HCA，140-token prefill | 对 Transformers reference 最大绝对误差 `1.79e-7` | +| Sliding decode | 最大绝对误差 `5.96e-8` | +| HCA decode | 最大绝对误差 `7.45e-8` | +| 正式 dense FP8 GEMM | 输出 finite；对显式 dequant reference 最大绝对误差 `0.142578`、平均绝对误差 `0.03160` | +| 正式 MXFP4 expert fused MoE | 输出 finite；对显式 MXFP4 dequant reference 最大绝对误差 `0.0009765625`、平均绝对误差 `0.00014795` | +| CUDA Graph 静态 CSA cache | 物理 GPU 6 捕获成功并连续 replay | +| CUDA Graph 内 EP+MXFP4 MoE | 物理 GPU 6,7；all-gather/fused-MoE/reduce-scatter replay 对 eager 最大误差 `0` | +| 完整 checkpoint index | 72,317 tensor；0 mapping error、0 unexpected skip | +| 正式 EP=4 rank-0 分片加载 | 48 shard、9,455 weight、42.50 GiB tensor，完整校验通过 | +| DeepSeek V4 及相关算子/并行/加载测试 | `140 passed, 1 skipped` | +| rebase 后全仓 CPU 测试 | `1357 passed, 138 skipped, 205 subtests passed` | + +全仓结果使用 `SPARSEVLLM_PLATFORM=cpu CUDA_VISIBLE_DEVICES=''`，skip 来自目标分支中需要 GPU 或可选依赖的环境测试。另使用真实 checkpoint 构造 `Config`，确认 runtime class 为 `DeepseekV4ForCausalLM`、cache class 为 `DeepseekV4CacheManager`、拓扑为 `DPA_EP`、cache 层数为 43，且正式 FP8/MXFP4 量化保持启用。 + +## 架构 rebase 回归 + +本实现已 rebase 到 `codex/qwen36-moe-sparse-methods` 的 `187c69b`，并按该分支的 `ModelSpec`、`RuntimeLayout`、`ParallelTopology` 和集中 checkpoint validator 重构。rebase 后完成上述专项测试、全仓 CPU 测试和真实 checkpoint 配置链验证。 + +下述正式 GPU 性能数据采集于本次架构 rebase 之前。rebase 收尾时物理 GPU 4–7 均在被其他任务使用，因此按照设备隔离要求没有抢占重跑；这些数字保留为同一模型实现的已有基线，不宣称为 rebase 后重新测得的数据。 + +## Tiny-random engine smoke + +以下数字只用于证明真实 engine 的 model construction、prefill、decode 和多进程 DPA 路径可运行，不代表正式模型性能。 + +| GPU | DP/EP | Prompt | Batch | Output | Prefill | Decode | TTFT | ITL | 单卡显存 | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 物理 GPU 6 | 1/1 | 32 | 1 | 2 | `1140.8 tok/s` | `39.3 tok/s` | `0.03 s` | `25.44 ms` | `3.41 GB` | +| 物理 GPU 6,7 | 2/2 | 32 | 3 | 2 | `1208.8 tok/s` | `57.8 tok/s` | `0.08 s` | `51.93 ms` | `3.41 GB` | + +双卡用 batch 3 故意制造不均衡 owner 数量，以覆盖 rank 补齐和全局 logits 重排。 + +## 正式 rank 分片加载 + +在空闲物理 GPU 6 上构造 EP=4 的 rank 0 分片并读取全部 48 个 shard；该检查不执行缺失其他 EP rank 的 forward。 + +| 指标 | 结果 | +| --- | ---: | +| 模型构造 | `2.02 s` | +| 权重加载 | `39.17 s` | +| 本地 weight 数 | `9,455` | +| loader tensor bytes | `42.50 GiB` | +| 常驻 GPU allocation | `42.44 GiB` | +| 43 层 MXFP4 interleave | `0.98 s` | +| 峰值 GPU allocation | `43.24 GiB` | + +## 正式模型推理与性能 + +### 环境和方法 + +- 设备：空闲物理 GPU 6,7，`NVIDIA H100 80GB HBM3`，driver `580.65.06`。 +- 软件：PyTorch `2.11.0+cu130`、CUDA runtime `13.0`、Transformers `5.13.1`、FlashInfer `0.6.15.post1`。 +- checkpoint：48/48 shard；`TP=1`、`DP=EP=2`；无 MTP、无可选稀疏方法。 +- engine：`max_model_len=4`、batch 1、prompt 1 token、output 2 token、`gpu_memory_utilization=0.995`。 +- graph：`decode_cuda_graph=true`，capture batch `[1]`、context `[4]`，sampling 不进入 graph。 +- 性能口径：engine warmup 后先完成一次请求，再连续执行 5 次相同请求；下表为这 5 次的中位数。Prefill/decode step 时间均包含本进程调度、模型执行和采样开销。 + +可复现实验的关键 kwargs： + +```python +llm = LLM( + "/data1/gqs/models/DeepSeek-V4-Flash-0731", + tensor_parallel_size=1, + data_parallel_size=2, + expert_parallel_size=2, + decode_cuda_graph=True, + decode_cuda_graph_capture_sizes=[1], + decode_cuda_graph_context_sizes=[4], + enforce_eager=False, + max_model_len=4, + max_num_batched_tokens=4, + engine_prefill_chunk_size=4, + max_num_seqs_in_batch=1, + max_decoding_seqs=1, + max_num_seqs_in_gpu=1, + mlp_chunk_size=4, + gpu_memory_utilization=0.995, + weight_loading_workers=2, +) +``` + +### 结果 + +| 指标 | 结果 | +| --- | ---: | +| 完整 engine 初始化 | `35.43 s` | +| 两 rank 权重加载 | `24.53 s` / `24.86 s` | +| 每 rank loader tensor bytes | `76.77 GiB` | +| 首个实测请求 | `444.6 ms` | +| 5 次请求总时间 | `2.2189 s` | +| 请求 E2E P50 | `443.07 ms` | +| Prefill step P50 | `390.78 ms` / `2.56 tok/s` | +| Decode step P50 / ITL | `52.31 ms` / `19.12 tok/s` | +| rank 0 PyTorch 常驻 / 峰值 allocation | `76.78 / 76.81 GiB` | +| 生成 token | `[294, 201]` | +| 捕获 graph 数 | `1` | +| graph key | `method='', batch_size=1, context_capacity=4, is_long_text=False, capture_sampling=False` | + +完整权重、cache、DPA all-gather、MXFP4 EP fused MoE、reduce-scatter、全局 logits 重排均经过真实 engine forward。后续 graph replay 成功返回且两 rank 正常退出；若任一 rank 的 graph 或 collective 序列不一致，该路径会同步报错或超时。 + +以上是双卡极限显存下的首版正确性基线，不代表推荐服务吞吐。`DP=EP=4` 的正常上下文性能尚未测量，因为本次验收期间物理 GPU 4,5 被其他任务持续占用，按照设备隔离要求未抢占；四卡数据应单独补测，不能和本表的短上下文结果直接比较。 diff --git a/docs/zh/design/README.md b/docs/zh/design/README.md index ac2c815a..0e52aa79 100644 --- a/docs/zh/design/README.md +++ b/docs/zh/design/README.md @@ -4,3 +4,4 @@ - [架构](architecture.md) - [Sparse-vLLM 控制图](control-map.md) +- [DeepSeek V4 Flash 运行时](deepseek-v4-flash.md) diff --git a/docs/zh/design/deepseek-v4-flash.md b/docs/zh/design/deepseek-v4-flash.md new file mode 100644 index 00000000..815f816c --- /dev/null +++ b/docs/zh/design/deepseek-v4-flash.md @@ -0,0 +1,92 @@ +# DeepSeek V4 Flash 运行时 + +本文说明 `/data1/gqs/models/DeepSeek-V4-Flash-0731` 在 Sparse-vLLM 中的首版实现。目标是正确运行 CUDA Graph decode 和 DeepSeek 风格 DPA+EP；首版不执行 checkpoint 中的 MTP，也不叠加 Sparse-vLLM 可选稀疏方法。 + +## 架构落点 + +实现已对齐 `codex/qwen36-moe-sparse-methods` 的模型、布局和并行抽象： + +- `ModelSpec` 声明模型运行类、专用 cache manager、DPA+EP 模式，以及 tiny-random 的量化和 checkpoint 校验策略。 +- `models/checkpoint.py` 集中校验正式 checkpoint 的架构、维度和 FP8/MXFP4 格式；模型实现不重复承担配置识别。 +- `RuntimeLayout` 将 CSA/HCA 识别为持有 KV 状态的完整 attention 层，统一生成 43 层 cache 布局。 +- `ParallelTopology.DPA_EP` 定义 `TP=1`、`DP=EP`、attention singleton group，以及重叠的 data/expert group；`ParallelContext` 只提供 collective。 +- runtime compatibility 由 `(model_type, parallel_mode)` 注册，DeepSeek V4 首版只允许 vanilla 方法。 +- `model_runner.py` 只保留执行期的 owner 分片、固定 shape 补齐、全局 logits 重排和 CUDA Graph 约束。 + +## 支持边界 + +- checkpoint 架构必须为 `DeepseekV4ForCausalLM`，dense 权重为 E4M3 FP8、动态激活量化、UE8M0 `128 x 128` scale，expert 权重为 MXFP4、K32 UE8M0 scale。 +- 运行布局固定为 `TP=1`、`DP=EP`，expert 数必须能被 EP 整除。 +- Hopper 首版要求 SM90、CUDA runtime 12.8 或更新版本，以及仓库环境中的 FlashInfer。 +- 仅支持 `vllm_sparse_method=""`。MTP、prefix cache 和其他可选稀疏方法会在配置阶段显式拒绝。 +- decode attention 按稳定的 `seq_id % DP` owner 分片；MoE 在重叠的全局 EP group 中执行。首版 prefill 在各 DPA rank 复制执行，以保证所有 rank 都有完整、可校验的初始 cache。 + +## 模型和算子 + +原生实现位于 `src/sparsevllm/models/deepseek_v4_native.py`： + +- mHC attention/FFN connection 和 hyper head 保留 fp32 Sinkhorn 与归一化语义。 +- Q/KV/O、shared expert 等 dense projection 复用 Sparse-vLLM 的 FP8 linear provider。 +- routed expert 仅在所属 EP rank 分配。checkpoint 的 `w1/w3/w2` packed MXFP4 bit pattern 原样加载，再由 FlashInfer Hopper W4A16 fused MoE 完成 interleave 和执行。 +- Router 支持前三层 hash routing 和其余层 `sqrt(softplus(x))` routing；SwiGLU clamp 固定使用 checkpoint 的 limit 10。 +- MTP tensor 只被识别并跳过，不进入参数分配和 forward。 + +完整 safetensors index 包含 72,317 个 tensor。EP=4、rank 0 的全索引映射审计结果为 0 个未知跳过、0 个目标缺失；远端 expert 和 MTP 是仅有的有意跳过项。 + +## Cache 和 CUDA Graph + +`DeepseekV4CacheManager` 按请求行号持有以下状态： + +- 每层一个 128-token shared-KV 环形窗口； +- 21 个 CSA 层每 4 token 一个 512-dim compressed KV 和 128-dim index KV； +- 20 个 HCA 层每 128 token 一个 512-dim compressed KV； +- CSA 的前窗 Ca overlap，以及 CSA/HCA 当前压缩窗口的 KV/gate ring。 + +decode 只通过 tensor index 原地更新这些状态，不创建或修改 Python `DynamicCache`，因此同一组地址可被 CUDA Graph 捕获和 replay。长 prefill 只把最后一个不重复的 sliding window 写回 ring，避免 advanced indexing 对重复 ring column 的未定义覆盖顺序。 + +Cache 按 `max_model_len` 和 `max_num_seqs_in_gpu` 预分配。若估算值超过 `gpu_memory_utilization` 给出的预算，启动会直接报出所需和可用 GiB；不会静默缩小上下文或请求容量。 + +## DPA+EP decode + +每个 decode step 的流程为： + +1. 所有 rank 收到同一全局请求列表，并选出本 rank 稳定 owner 的请求。 +2. 请求较少的 rank 使用非 owner cache 行补齐到相同 shape；补齐输出会被丢弃。 +3. attention 在本地 singleton group 执行。 +4. 每个 MoE 层 all-gather hidden state 和 token id，各 EP rank 只计算本地 expert，再 reduce-scatter 回 attention rank；shared expert 保持本地执行。 +5. full-vocab logits 在 EP group 中聚合，rank 0 按原始请求顺序重排并统一采样。 + +固定补齐 shape 使所有 rank 选择相同的 CUDA Graph batch family，也使 graph 内 NCCL collective 的调用顺序一致。 + +## 推荐启动参数 + +正式四卡短上下文服务可从以下 engine kwargs 开始，再按目标上下文增加 `max_model_len`： + +```json +{ + "tensor_parallel_size": 1, + "data_parallel_size": 4, + "expert_parallel_size": 4, + "decode_cuda_graph": true, + "enforce_eager": false, + "max_model_len": 512, + "max_num_seqs_in_batch": 4, + "max_decoding_seqs": 4, + "max_num_seqs_in_gpu": 4, + "max_num_batched_tokens": 512, + "engine_prefill_chunk_size": 512, + "gpu_memory_utilization": 0.9 +} +``` + +当前已在两张 H100 80GB 上以 `DP=EP=2` 完成完整 checkpoint 推理和 CUDA Graph replay。该布局每 rank 常驻约 76.78 GiB，只适合 `max_model_len=4` 的最小正确性与性能验收；正常上下文应使用 `DP=EP=4` 或更多 rank，不能依靠静默缩减 cache 容量。 + +开发阶段设置： + +```bash +export SPARSEVLLM_TINY_RANDOM=1 +export SPARSEVLLM_TINY_RANDOM_CONFIG="$PWD/configs/debug/deepseek_v4_flash_tiny_random.json" +export SPARSEVLLM_TINY_RANDOM_SEED=17 +``` + +Tiny random 不读取 safetensors，且性能与生成质量不能代表正式模型。 diff --git a/docs/zh/features/supported-models.md b/docs/zh/features/supported-models.md index d273ac8d..9b67d2ee 100644 --- a/docs/zh/features/supported-models.md +++ b/docs/zh/features/supported-models.md @@ -12,8 +12,10 @@ | Qwen3 Dense | `qwen3` | BF16 / FP16 / 块级 FP8 | ✅（FP8：1/2/4/8） | 仅支持 1 | 仅支持 1 | | Qwen3MoE | `qwen3_moe` | BF16 / FP16 / 块级 FP8 | ✅（TP > 1 时模型 dtype 仅支持 BF16） | 仅支持 1 | ✅ | | Qwen3.5 / Qwen3.6 | `qwen3_5` | BF16 / 块级 FP8 | ✅ | 仅支持 1 | 仅支持 1 | +| Qwen3.6 MoE | `qwen3_5_moe` | BF16 / 块级 FP8 | ✅ | 仅支持 1 | ✅ | | Llama 3 / 3.1 | `llama` | BF16 / FP16 | ✅ | 仅支持 1 | 仅支持 1 | | MiniMax M2.7 | `minimax_m2` | 块级 FP8，非量化权重使用 BF16 | ✅ | 仅支持 1 | ✅ | +| DeepSeek V4 Flash | `deepseek_v4` | Dense E4M3 FP8 + Expert MXFP4 | 仅支持 1 | ✅，且 DP=EP | ✅，且 EP=DP | TP 规模限制为 1 到 8，并且 checkpoint 维度（包括 attention head 数和 vocabulary 大小）必须能被所选 TP 规模整除。Qwen3MoE 的 EP 规模必须整除 `num_experts`；MiniMax M2.7 的 EP 规模必须整除 `num_local_experts`。 @@ -24,7 +26,9 @@ size 为 `T`。该布局要求 `DP=1` 且 `T % E == 0`；专家数量必须能 TP 要求模型 dtype 为 BF16；FP16 Qwen3MoE checkpoint 仅支持 `TP=1`。当 `TP=1` 时，原有 EP 布局的 world size 为 `E`。 -块级 FP8 要求使用 E4M3 权重、动态激活量化以及 `128 x 128` 的权重块大小。Qwen3.5 和 Qwen3.6 的配置在内部统一规范为 `model_type=qwen3_5`。 +块级 FP8 要求使用 E4M3 权重、动态激活量化以及 `128 x 128` 的权重块大小。 +Qwen3.5/Qwen3.6 Dense 配置在内部统一规范为 `model_type=qwen3_5`；Qwen3.6 +MoE 使用 `model_type=qwen3_5_moe`。 ## 稀疏方法支持 @@ -34,8 +38,10 @@ TP 要求模型 dtype 为 BF16；FP16 Qwen3MoE checkpoint 仅支持 `TP=1`。当 | Qwen3 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | 需要压缩器² | | Qwen3MoE | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Qwen3.5 / Qwen3.6 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | 匹配的 checkpoint³ | +| Qwen3.6 MoE | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Llama 3 / 3.1 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | 指定 checkpoint¹ | 需要 compressor² | | MiniMax M2.7 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | +| DeepSeek V4 Flash | ✅ | — | — | — | — | — | — | — | — | — | ¹ SkipKV 仅支持已发布 steering vector 的模型： `DeepSeek-R1-Distill-Qwen-7B`、`DeepSeek-R1-Distill-Qwen-14B` 和 diff --git a/src/sparsevllm/configs/common.py b/src/sparsevllm/configs/common.py index 93d9ef0a..235a39a1 100644 --- a/src/sparsevllm/configs/common.py +++ b/src/sparsevllm/configs/common.py @@ -104,7 +104,3 @@ def _resolve_long_prefill_offload_threshold(configured: Any) -> int: if resolved is None: raise ValueError("long_prefill_offload_threshold must be a positive integer.") return int(resolved) - - -def _model_path_basename(model_path: str) -> str: - return str(model_path).rstrip("/").split("/")[-1] diff --git a/src/sparsevllm/configs/delta.py b/src/sparsevllm/configs/delta.py index 8751240e..f2ffc820 100644 --- a/src/sparsevllm/configs/delta.py +++ b/src/sparsevllm/configs/delta.py @@ -1,6 +1,8 @@ """DeltaKV configuration normalization and runtime validation.""" import importlib.util +import json +import os from typing import Any from sparsevllm.configs.common import ( @@ -9,10 +11,6 @@ _normalize_positive_int, _normalize_positive_multiple, ) -from sparsevllm.configs.model import ( - _is_qwen35_deltakv_checkpoint, - _qwen35_deltakv_message, -) from sparsevllm.utils.log import log_once def _flash_attn_available() -> bool: @@ -42,15 +40,6 @@ def _resolve_deltakv_sparse_decode_backend(value: Any) -> str: return backend -SUPPORTED_SKIPKV_MODEL_NAMES = frozenset( - { - "DeepSeek-R1-Distill-Llama-8B", - "DeepSeek-R1-Distill-Qwen-7B", - "DeepSeek-R1-Distill-Qwen-14B", - } -) - - def _normalize_deltakv_kernel_options(config) -> None: if int(config.deltakv_cluster_gather_chunk_size) <= 0: raise ValueError( @@ -150,7 +139,24 @@ def normalize_deltakv_storage(config) -> None: _normalize_deltakv_capacity(config) -def validate_deltakv_runtime(config, *, is_qwen35: bool) -> None: +def _checkpoint_targets_model(path: str | None, model_types: frozenset[str]) -> bool: + config_path = os.path.join(path, "config.json") if path and os.path.isdir(path) else None + if config_path is None or not os.path.isfile(config_path): + return False + with open(config_path, "r", encoding="utf-8") as f: + checkpoint_config = json.load(f) + return any( + str(checkpoint_config.get(field, "")).strip().lower() in model_types + for field in ( + "model_type", + "base_model_type", + "target_model_type", + "runtime_model_type", + ) + ) + + +def validate_deltakv_runtime(config) -> None: # Normalize compressor type strings. for attr in ("compressor_down_type", "compressor_up_type"): v = getattr(config, attr, "auto") @@ -165,8 +171,15 @@ def validate_deltakv_runtime(config, *, is_qwen35: bool) -> None: "verify results carefully before treating them as final.", level="WARNING", ) - if is_qwen35 and not _is_qwen35_deltakv_checkpoint(config.deltakv_path): - raise ValueError(_qwen35_deltakv_message()) + checkpoint_model_types = config.model_spec.deltakv_checkpoint_model_types + if checkpoint_model_types and not _checkpoint_targets_model( + config.deltakv_path, + checkpoint_model_types, + ): + raise ValueError( + f"DeltaKV for {config.model_spec.name} requires a compatible " + "deltakv_path. Use vllm_sparse_method='' to run vanilla inference." + ) if not bool(getattr(config, "use_compression", True)): raise ValueError("DeltaKV runtime is compressor-only; set use_compression=True.") if bool(getattr(config, "enable_sparse_ref_fp8", False)): diff --git a/src/sparsevllm/configs/model.py b/src/sparsevllm/configs/model.py index 0200316d..e808ac7e 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -1,550 +1,60 @@ -"""Model metadata loading, layout construction, and checkpoint validation.""" +"""Model metadata loading and validation orchestration.""" import json import os -from dataclasses import dataclass from types import SimpleNamespace from typing import Any -import torch from transformers import AutoConfig from sparsevllm.method_registry import ( - H2O_SUPPORTED_MODEL_TYPES, validate_model_runtime_compatibility, + validate_sparse_method_assets, ) +from sparsevllm.models.checkpoint import validate_checkpoint +from sparsevllm.models.layout import RuntimeLayout +from sparsevllm.models.spec import ModelSpec, resolve_model_spec +from sparsevllm.quantization import QuantizationConfig +from sparsevllm.utils.config import config_get from sparsevllm.utils.log import logger, log_once -try: - from transformers import Qwen3Config -except ImportError: - Qwen3Config = AutoConfig - -def _config_get(config: Any, name: str, default: Any = None) -> Any: - if config is None: - return default - if isinstance(config, dict): - return config.get(name, default) - return getattr(config, name, default) - def _config_to_namespace(config: dict[str, Any]) -> SimpleNamespace: return SimpleNamespace(**config) -def _load_raw_qwen35_config(model_path: str, error: Exception) -> SimpleNamespace: +def _load_model_config(model_path: str) -> Any: + try: + return AutoConfig.from_pretrained(model_path, trust_remote_code=True) + except Exception as error: + load_error = error config_path = os.path.join(model_path, "config.json") if not os.path.isfile(config_path): raise RuntimeError( - "AutoConfig.from_pretrained failed and no config.json exists for explicit " - f"qwen3_5 fallback. model={model_path} error={type(error).__name__}: {error}" - ) from error + "AutoConfig.from_pretrained failed and no config.json exists for an " + f"explicit raw-config fallback. model={model_path} " + f"error={type(load_error).__name__}: {load_error}" + ) from load_error with open(config_path, "r", encoding="utf-8") as f: raw_config = json.load(f) - if not _is_qwen35_outer_config(raw_config): + model_type = str(config_get(raw_config, "model_type", "") or "") + model_spec = resolve_model_spec(model_type) + if not model_spec.allow_raw_config: raise RuntimeError( "AutoConfig.from_pretrained failed. Refusing to silently fall back to raw " - f"`config.json` for non-qwen3_5 model. model={model_path} " - f"error={type(error).__name__}: {error}" - ) from error + f"config.json for {model_spec.name}. model={model_path} " + f"error={type(load_error).__name__}: {load_error}" + ) from load_error log_once( - "AutoConfig.from_pretrained failed for qwen3_5/qwen3_6; loading raw config.json " - "through Sparse-vLLM's explicit mixed-runtime parser.", + f"AutoConfig.from_pretrained failed for {model_spec.name}; loading raw " + "config.json through its explicit model specification.", level="WARNING", ) return _config_to_namespace(raw_config) -def _coerce_int_list(name: str, value: Any, *, allow_none: bool = False) -> list[int] | None: - if value is None: - if allow_none: - return None - raise ValueError(f"{name} is required.") - if isinstance(value, str): - raw = value.strip() - if not raw: - return [] - parts = [part.strip() for part in raw.split(",") if part.strip()] - return [int(part) for part in parts] - if isinstance(value, (list, tuple)): - return [int(item) for item in value] - raise ValueError(f"{name} must be a list/tuple of ints or a comma-separated string, got {value!r}.") - - -def _attention_type_is_full(value: Any) -> bool: - text = str(value).strip().lower() - return text in {"full", "full_attention", "attention", "self_attention", "sliding_attention"} - - -def _attention_type_is_linear(value: Any) -> bool: - text = str(value).strip().lower() - return text in {"linear", "linear_attention", "recurrent", "recurrent_attention", "gated_delta", "gated_delta_net"} - - -@dataclass(frozen=True) -class QuantizationConfig: - enabled: bool = False - quant_method: str = "" - weight_dtype: str = "" - activation_scheme: str = "" - weight_block_size: tuple[int, int] | None = None - model_name: str = "qwen3_5" - - @classmethod - def disabled(cls) -> "QuantizationConfig": - return cls() - - def to_dict(self) -> dict[str, Any]: - if not self.enabled: - return {} - payload: dict[str, Any] = { - "quant_method": self.quant_method, - "fmt": self.weight_dtype, - "activation_scheme": self.activation_scheme, - } - if self.weight_block_size is not None: - payload["weight_block_size"] = list(self.weight_block_size) - return payload - - @classmethod - def from_hf_config( - cls, - value: Any, - *, - required_fp8: bool = False, - model_name: str = "qwen3_5", - ) -> "QuantizationConfig": - if value is None: - if required_fp8: - raise ValueError( - f"{model_name} requires FP8 quantization_config; " - "BF16/FP16 fallback is not supported." - ) - return cls.disabled() - - quant_method = str( - _config_get(value, "quant_method", _config_get(value, "method", "")) - or "" - ).strip().lower() - if quant_method not in {"fp8", "fbgemm_fp8"}: - if required_fp8: - raise ValueError( - f"{model_name} requires quantization_config.quant_method='fp8', " - f"got {quant_method!r}." - ) - return cls.disabled() - - weight_dtype = str( - _config_get( - value, - "weight_dtype", - _config_get(value, "fmt", _config_get(value, "format", "e4m3")), - ) - or "" - ).strip().lower() - if "e4m3" not in weight_dtype: - raise ValueError( - f"Sparse-vLLM {model_name} FP8 supports e4m3 weights only, " - f"got weight_dtype={weight_dtype!r}." - ) - - activation_scheme = str( - _config_get(value, "activation_scheme", _config_get(value, "activation", "dynamic")) - or "" - ).strip().lower() - if activation_scheme != "dynamic": - raise ValueError( - f"Sparse-vLLM {model_name} FP8 supports dynamic activation only, " - f"got activation_scheme={activation_scheme!r}." - ) - - block_size = _config_get( - value, - "weight_block_size", - _config_get(value, "weight_block_shape", _config_get(value, "block_size", (128, 128))), - ) - if isinstance(block_size, int): - block_tuple = (int(block_size), int(block_size)) - elif isinstance(block_size, (list, tuple)) and len(block_size) == 2: - block_tuple = (int(block_size[0]), int(block_size[1])) - else: - raise ValueError(f"weight_block_size must be a pair, got {block_size!r}.") - if block_tuple != (128, 128): - raise ValueError( - f"Sparse-vLLM {model_name} FP8 supports " - "weight_block_size=(128, 128) only, " - f"got {block_tuple}." - ) - - return cls( - enabled=True, - quant_method="fp8", - weight_dtype="e4m3", - activation_scheme="dynamic", - weight_block_size=block_tuple, - model_name=model_name, - ) - - -def _validate_qwen35_checkpoint_precision( - hf_config: Any, - raw_quantization_config: Any, - quantization_config: QuantizationConfig, -) -> None: - if quantization_config.enabled: - return - - quant_method = str( - _config_get( - raw_quantization_config, - "quant_method", - _config_get(raw_quantization_config, "method", ""), - ) - or "" - ).strip().lower() - if quant_method: - raise NotImplementedError( - "qwen3_5 supports unquantized BF16 or block FP8 checkpoints only, " - f"got quant_method={quant_method!r}." - ) - - configured_dtype = _config_get(hf_config, "torch_dtype", None) - if configured_dtype not in {torch.bfloat16, "bfloat16"}: - raise NotImplementedError( - "Unquantized qwen3_5 checkpoints require BF16 weights, " - f"got torch_dtype={configured_dtype!r}." - ) - - -_MINIMAX_M2_FIXED_FIELDS = { - "vocab_size": 200064, - "hidden_size": 3072, - "intermediate_size": 1536, - "num_hidden_layers": 62, - "num_attention_heads": 48, - "num_key_value_heads": 8, - "head_dim": 128, - "rotary_dim": 64, - "num_local_experts": 256, - "num_experts_per_tok": 8, - "max_position_embeddings": 204800, - "shared_intermediate_size": 0, - "mtp_transformer_layers": 1, - "num_mtp_modules": 3, -} - - -def _validate_minimax_m2_checkpoint_config( - hf_config: Any, - raw_quantization_config: Any, -) -> None: - architectures = tuple(_config_get(hf_config, "architectures", ()) or ()) - if architectures != ("MiniMaxM2ForCausalLM",): - raise ValueError( - "MiniMax M2.7 requires architectures=['MiniMaxM2ForCausalLM'], " - f"got {list(architectures)}." - ) - for field_name, expected in _MINIMAX_M2_FIXED_FIELDS.items(): - actual = _config_get(hf_config, field_name, None) - if actual != expected: - raise ValueError( - f"MiniMax M2.7 requires {field_name}={expected!r}, got {actual!r}." - ) - - expected_values = { - "hidden_act": "silu", - "qk_norm_type": "per_layer", - "scoring_func": "sigmoid", - "use_qk_norm": True, - "use_routing_bias": True, - "use_mtp": True, - "tie_word_embeddings": False, - } - for field_name, expected in expected_values.items(): - actual = _config_get(hf_config, field_name, None) - if actual != expected: - raise ValueError( - f"MiniMax M2.7 requires {field_name}={expected!r}, got {actual!r}." - ) - - configured_dtype = _config_get(hf_config, "torch_dtype", None) - if configured_dtype is None: - configured_dtype = _config_get(hf_config, "dtype", None) - if configured_dtype not in {torch.bfloat16, "bfloat16"}: - raise ValueError( - "MiniMax M2.7 requires BF16 non-quantized parameters, " - f"got dtype={configured_dtype!r}." - ) - - excluded_modules = { - str(name) - for name in ( - _config_get(raw_quantization_config, "modules_to_not_convert", ()) or () - ) - } - required_exclusions = {"gate", "e_score_correction_bias", "lm_head"} - missing_exclusions = sorted(required_exclusions - excluded_modules) - if missing_exclusions: - raise ValueError( - "MiniMax M2.7 quantization_config must exclude gate, " - "e_score_correction_bias, and lm_head; missing " - f"{missing_exclusions}." - ) - - -def _validate_qwen3_moe_fp8_checkpoint_config( - hf_config: Any, - raw_quantization_config: Any, -) -> None: - architectures = tuple(_config_get(hf_config, "architectures", ()) or ()) - if architectures != ("Qwen3MoeForCausalLM",): - raise ValueError( - "Qwen3MoE FP8 requires architectures=['Qwen3MoeForCausalLM'], " - f"got {list(architectures)}." - ) - configured_dtype = _config_get(hf_config, "torch_dtype", None) - if configured_dtype is None: - configured_dtype = _config_get(hf_config, "dtype", None) - if configured_dtype not in {torch.bfloat16, "bfloat16"}: - raise ValueError( - "Qwen3MoE FP8 requires BF16 non-quantized parameters, " - f"got dtype={configured_dtype!r}." - ) - - hidden_size = int(_config_get(hf_config, "hidden_size", 0) or 0) - intermediate_size = int( - _config_get(hf_config, "moe_intermediate_size", 0) or 0 - ) - if hidden_size % 128 or intermediate_size % 128: - raise ValueError( - "Qwen3MoE FP8 requires hidden_size and moe_intermediate_size " - f"aligned to 128, got {hidden_size}/{intermediate_size}." - ) - - excluded_modules = { - str(name) - for name in ( - _config_get(raw_quantization_config, "modules_to_not_convert", ()) or () - ) - } - num_layers = int(_config_get(hf_config, "num_hidden_layers", 0) or 0) - required_exclusions = {"lm_head"} - required_exclusions.update( - f"model.layers.{layer_idx}.mlp.gate" - for layer_idx in range(num_layers) - ) - missing_exclusions = sorted(required_exclusions - excluded_modules) - if missing_exclusions: - raise ValueError( - "Qwen3MoE FP8 quantization_config must exclude lm_head and every " - f"router gate; missing {missing_exclusions[:8]}." - ) - - -def _validate_qwen3_fp8_checkpoint_config( - hf_config: Any, - *, - tensor_parallel_size: int, -) -> None: - architectures = tuple(_config_get(hf_config, "architectures", ()) or ()) - if architectures != ("Qwen3ForCausalLM",): - raise ValueError( - "Qwen3 FP8 requires architectures=['Qwen3ForCausalLM'], " - f"got {list(architectures)}." - ) - configured_dtype = _config_get(hf_config, "torch_dtype", None) - if configured_dtype is None: - configured_dtype = _config_get(hf_config, "dtype", None) - if configured_dtype not in {torch.bfloat16, "bfloat16"}: - raise ValueError( - "Qwen3 FP8 requires BF16 non-quantized parameters, " - f"got dtype={configured_dtype!r}." - ) - - tp_size = int(tensor_parallel_size) - head_dim = int(_config_get(hf_config, "head_dim", 0) or 0) - dimensions = { - "hidden_size": int(_config_get(hf_config, "hidden_size", 0) or 0), - "intermediate_size": int( - _config_get(hf_config, "intermediate_size", 0) or 0 - ), - "query_size": int( - _config_get(hf_config, "num_attention_heads", 0) or 0 - ) - * head_dim, - "key_value_size": int( - _config_get(hf_config, "num_key_value_heads", 0) or 0 - ) - * head_dim, - } - invalid_dimensions = { - name: size - for name, size in dimensions.items() - if size <= 0 or size % (128 * tp_size) - } - if invalid_dimensions: - raise ValueError( - "Qwen3 FP8 requires every TP-local dense projection dimension to be " - "128-aligned; " - f"TP={tp_size}, invalid={invalid_dimensions}." - ) - - -@dataclass(frozen=True) -class RuntimeLayout: - num_layers: int - num_kv_layers: int - full_attention_layer_indices: tuple[int, ...] - linear_attention_layer_indices: tuple[int, ...] - layer_idx_to_kv_idx: tuple[int | None, ...] - kv_idx_to_layer_idx: tuple[int, ...] - - @classmethod - def dense(cls, num_layers: int) -> "RuntimeLayout": - num_layers = int(num_layers) - layers = tuple(range(num_layers)) - return cls( - num_layers=num_layers, - num_kv_layers=num_layers, - full_attention_layer_indices=layers, - linear_attention_layer_indices=(), - layer_idx_to_kv_idx=tuple(range(num_layers)), - kv_idx_to_layer_idx=layers, - ) - - @classmethod - def from_config(cls, hf_config: Any, *, require_mixed: bool = False) -> "RuntimeLayout": - num_layers = int(_config_get(hf_config, "num_hidden_layers")) - layer_types = _config_get(hf_config, "layer_types", None) - full_layers = _coerce_int_list( - "full_attention_layer_indices", - _config_get( - hf_config, - "full_attention_layer_indices", - _config_get(hf_config, "attention_layer_indices", None), - ), - allow_none=True, - ) - linear_layers = _coerce_int_list( - "linear_attention_layer_indices", - _config_get(hf_config, "linear_attention_layer_indices", None), - allow_none=True, - ) - - if layer_types is not None: - if len(layer_types) != num_layers: - raise ValueError( - f"runtime layer_types length must equal num_hidden_layers: " - f"{len(layer_types)} != {num_layers}." - ) - inferred_full: list[int] = [] - inferred_linear: list[int] = [] - for idx, layer_type in enumerate(layer_types): - if _attention_type_is_full(layer_type): - inferred_full.append(idx) - elif _attention_type_is_linear(layer_type): - inferred_linear.append(idx) - else: - raise ValueError(f"Unsupported qwen3_5 layer_types[{idx}]={layer_type!r}.") - full_layers = inferred_full if full_layers is None else full_layers - linear_layers = inferred_linear if linear_layers is None else linear_layers - - if full_layers is None and linear_layers is None: - if require_mixed: - raise ValueError( - "qwen3_5 requires a mixed attention layer map: provide layer_types or " - "full_attention_layer_indices/linear_attention_layer_indices." - ) - return cls.dense(num_layers) - if full_layers is None: - linear_set = set(linear_layers or []) - full_layers = [idx for idx in range(num_layers) if idx not in linear_set] - if linear_layers is None: - full_set = set(full_layers or []) - linear_layers = [idx for idx in range(num_layers) if idx not in full_set] - - full_tuple = tuple(sorted(int(idx) for idx in full_layers)) - linear_tuple = tuple(sorted(int(idx) for idx in linear_layers)) - full_set = set(full_tuple) - linear_set = set(linear_tuple) - expected = set(range(num_layers)) - if full_set & linear_set: - overlap = sorted(full_set & linear_set) - raise ValueError(f"RuntimeLayout full and linear layer sets overlap: {overlap}.") - if full_set | linear_set != expected: - missing = sorted(expected - (full_set | linear_set)) - extra = sorted((full_set | linear_set) - expected) - raise ValueError(f"RuntimeLayout layer map is incomplete: missing={missing}, extra={extra}.") - - raw_layer_to_kv = _config_get(hf_config, "layer_idx_to_kv_idx", None) - if raw_layer_to_kv is None: - layer_to_kv: list[int | None] = [None] * num_layers - for kv_idx, layer_idx in enumerate(full_tuple): - layer_to_kv[layer_idx] = kv_idx - else: - if len(raw_layer_to_kv) != num_layers: - raise ValueError( - "layer_idx_to_kv_idx length must equal num_hidden_layers: " - f"{len(raw_layer_to_kv)} != {num_layers}." - ) - layer_to_kv = [] - for idx, value in enumerate(raw_layer_to_kv): - if value is None or int(value) < 0: - layer_to_kv.append(None) - else: - layer_to_kv.append(int(value)) - for layer_idx in linear_tuple: - if layer_to_kv[layer_idx] is not None: - raise ValueError( - f"layer_idx_to_kv_idx[{layer_idx}] must be None/-1 for linear_attention layers." - ) - - kv_pairs = [(kv_idx, layer_idx) for layer_idx, kv_idx in enumerate(layer_to_kv) if kv_idx is not None] - if len(kv_pairs) != len(full_tuple): - raise ValueError( - "RuntimeLayout must assign exactly one KV index to each full_attention layer: " - f"full_layers={len(full_tuple)} assigned={len(kv_pairs)}." - ) - kv_pairs.sort() - kv_indices = [kv_idx for kv_idx, _ in kv_pairs] - if kv_indices != list(range(len(kv_pairs))): - raise ValueError(f"KV layer indices must be contiguous from 0, got {kv_indices}.") - kv_tuple = tuple(layer_idx for _, layer_idx in kv_pairs) - - configured_num_kv_layers = _config_get(hf_config, "num_kv_layers", None) - if configured_num_kv_layers is not None and int(configured_num_kv_layers) != len(kv_tuple): - raise ValueError( - f"num_kv_layers={configured_num_kv_layers} does not match full_attention layers={len(kv_tuple)}." - ) - return cls( - num_layers=num_layers, - num_kv_layers=len(kv_tuple), - full_attention_layer_indices=full_tuple, - linear_attention_layer_indices=linear_tuple, - layer_idx_to_kv_idx=tuple(layer_to_kv), - kv_idx_to_layer_idx=kv_tuple, - ) - - def is_full_attention(self, layer_idx: int) -> bool: - return self.layer_idx_to_kv_idx[int(layer_idx)] is not None - - def is_linear_attention(self, layer_idx: int) -> bool: - return self.layer_idx_to_kv_idx[int(layer_idx)] is None - - def kv_layer_index(self, layer_idx: int) -> int: - layer_idx = int(layer_idx) - kv_idx = self.layer_idx_to_kv_idx[layer_idx] - if kv_idx is None: - raise RuntimeError(f"layer_idx={layer_idx} is linear_attention and has no KV cache") - return int(kv_idx) - - -def _is_qwen35_outer_config(config: Any) -> bool: - return str(_config_get(config, "model_type", "") or "").strip().lower() in {"qwen3_5", "qwen3_6"} - - def _extract_text_config(config: Any) -> Any: - text_config = _config_get(config, "text_config", None) + text_config = config_get(config, "text_config", None) if text_config is None: return config if isinstance(text_config, dict): @@ -552,214 +62,11 @@ def _extract_text_config(config: Any) -> Any: return text_config -def _qwen35_deltakv_message() -> str: - return ( - "DeltaKV for qwen3_5 requires a qwen3_5-compatible deltakv_path. " - "Use vllm_sparse_method='' to run quantized vanilla inference." - ) - - -def _is_qwen35_deltakv_checkpoint(path: str | None) -> bool: - if path is None or not os.path.exists(path): - return False - config_path = os.path.join(path, "config.json") if os.path.isdir(path) else None - if config_path is None or not os.path.isfile(config_path): - return False - with open(config_path, "r", encoding="utf-8") as f: - checkpoint_config = json.load(f) - candidates = [ - checkpoint_config.get("model_type"), - checkpoint_config.get("base_model_type"), - checkpoint_config.get("target_model_type"), - checkpoint_config.get("runtime_model_type"), - ] - return any(str(value).strip().lower() in {"qwen3_5", "qwen3_6"} for value in candidates if value) - - - - -def _validate_runtime_compatibility(config, *, model_type: str) -> None: - validate_model_runtime_compatibility( - model_type=model_type, - sparse_method=config.vllm_sparse_method, - tensor_parallel_size=config.tensor_parallel_size, - expert_parallel_size=config.expert_parallel_size, - data_parallel_size=config.data_parallel_size, - enforce_eager=config.enforce_eager, - decode_cuda_graph=config.decode_cuda_graph, - enable_prefix_caching=config.enable_prefix_caching, - ) - - -def _validate_qwen3_moe_runtime(config, *, model_type: str) -> None: - tp_size = int(config.tensor_parallel_size) - ep_size = int(config.expert_parallel_size) - if config.data_parallel_size != 1: - raise ValueError( - "Qwen3MoE requires DP=1, got " - f"TP={config.tensor_parallel_size}, EP={config.expert_parallel_size}, " - f"DP={config.data_parallel_size}." - ) - if tp_size > 1 and tp_size % ep_size: - raise ValueError( - "Qwen3MoE outer tensor_parallel_size must be divisible by " - f"expert_parallel_size, got outer TP={tp_size}, MoE EP={ep_size}." - ) - num_experts = int(getattr(config.hf_config, "num_experts", 0) or 0) - if num_experts <= 0: - raise ValueError(f"Qwen3MoE requires a positive num_experts, got {num_experts}.") - if ep_size > num_experts: - raise ValueError( - "expert_parallel_size must not exceed num_experts, " - f"got EP={config.expert_parallel_size}, num_experts={num_experts}." - ) - if num_experts % ep_size != 0: - raise ValueError( - "Qwen3MoE requires num_experts divisible by expert_parallel_size, " - f"got num_experts={num_experts}, EP={config.expert_parallel_size}." - ) - if tp_size > 1: - divisible_fields = { - "num_attention_heads": int(config.hf_config.num_attention_heads), - "num_key_value_heads": int(config.hf_config.num_key_value_heads), - "vocab_size": int(config.hf_config.vocab_size), - } - for field, value in divisible_fields.items(): - if value % tp_size: - raise ValueError( - f"Qwen3MoE {field} must be divisible by tensor_parallel_size, " - f"got {value} and {tp_size}." - ) - moe_tp_size = tp_size // ep_size - moe_intermediate_size = int(config.hf_config.moe_intermediate_size) - if moe_intermediate_size % moe_tp_size: - raise ValueError( - "Qwen3MoE moe_intermediate_size must be divisible by MoE TP size, " - f"got {moe_intermediate_size} and {moe_tp_size}." - ) - top_k = int(getattr(config.hf_config, "num_experts_per_tok", 0) or 0) - if not 1 <= top_k <= num_experts: - raise ValueError( - "Qwen3MoE num_experts_per_tok must be in [1, num_experts], " - f"got top_k={top_k}, num_experts={num_experts}." - ) - decoder_sparse_step = int( - getattr(config.hf_config, "decoder_sparse_step", 1) - ) - mlp_only_layers = tuple( - int(layer_idx) - for layer_idx in (getattr(config.hf_config, "mlp_only_layers", ()) or ()) - ) - if decoder_sparse_step != 1 or mlp_only_layers: - raise NotImplementedError( - "Qwen3MoE v1 requires every decoder layer to be MoE, got " - f"decoder_sparse_step={decoder_sparse_step}, " - f"mlp_only_layers={list(mlp_only_layers)}." - ) - shared_intermediate_size = int( - getattr(config.hf_config, "shared_expert_intermediate_size", 0) or 0 +def _finalize_model_config(config, model_spec: ModelSpec) -> None: + config.runtime_layout = RuntimeLayout.from_config( + config.hf_config, + require_mixed=model_spec.mixed_attention, ) - if shared_intermediate_size != 0: - raise NotImplementedError( - "Qwen3MoE v1 does not support shared experts, got " - f"shared_expert_intermediate_size={shared_intermediate_size}." - ) - model_dtype = getattr(config.hf_config, "torch_dtype", None) - if model_dtype not in {torch.bfloat16, torch.float16}: - raise NotImplementedError( - "Qwen3MoE v1 supports BF16/FP16 checkpoints only, " - f"got torch_dtype={model_dtype}." - ) - if tp_size > 1 and model_dtype != torch.bfloat16: - raise NotImplementedError( - "Qwen3MoE outer TP supports BF16 checkpoints only, " - f"got torch_dtype={model_dtype}." - ) - _validate_runtime_compatibility(config, model_type=model_type) - - -def _validate_minimax_runtime(config, *, model_type: str) -> None: - tp_size = int(config.tensor_parallel_size) - ep_size = int(config.expert_parallel_size) - if config.data_parallel_size != 1: - raise ValueError( - "MiniMax M2.7 requires DP=1, got " - f"TP={config.tensor_parallel_size}, EP={config.expert_parallel_size}, " - f"DP={config.data_parallel_size}." - ) - if tp_size > 1 and tp_size % ep_size: - raise ValueError( - "MiniMax M2.7 outer tensor_parallel_size must be divisible by " - f"expert_parallel_size, got outer TP={tp_size}, MoE EP={ep_size}." - ) - num_experts = int(getattr(config.hf_config, "num_local_experts")) - if ep_size > num_experts: - raise ValueError( - "MiniMax M2.7 expert_parallel_size must not exceed " - f"num_local_experts={num_experts}, got {config.expert_parallel_size}." - ) - if num_experts % ep_size != 0: - raise ValueError( - "MiniMax M2.7 requires num_local_experts divisible by " - f"expert_parallel_size, got {num_experts} and " - f"{config.expert_parallel_size}." - ) - if tp_size > 1: - divisible_fields = { - "num_attention_heads": int(config.hf_config.num_attention_heads), - "num_key_value_heads": int(config.hf_config.num_key_value_heads), - "vocab_size": int(config.hf_config.vocab_size), - } - for field, value in divisible_fields.items(): - if value % tp_size: - raise ValueError( - f"MiniMax M2.7 {field} must be divisible by " - f"tensor_parallel_size, got {value} and {tp_size}." - ) - moe_tp_size = tp_size // ep_size - intermediate_size = int(config.hf_config.intermediate_size) - if intermediate_size % moe_tp_size: - raise ValueError( - "MiniMax M2.7 intermediate_size must be divisible by MoE TP size, " - f"got {intermediate_size} and {moe_tp_size}." - ) - _validate_runtime_compatibility(config, model_type=model_type) - - -def _validate_dense_parallelism(config, *, model_type: str) -> None: - if config.expert_parallel_size != 1 or config.data_parallel_size != 1: - raise ValueError( - f"Dense model_type={model_type!r} requires EP=1 and DP=1, got " - f"TP={config.tensor_parallel_size}, EP={config.expert_parallel_size}, " - f"DP={config.data_parallel_size}." - ) - - -def _finalize_model_config(config, *, is_qwen35: bool) -> None: - if ( - config.vllm_sparse_method == "deltakv" - and not is_qwen35 - and config.deltakv_path is None - and not config.allow_missing_deltakv_path - ): - raise ValueError( - "DeltaKV requires deltakv_path for compressor sparse layers. " - "Set allow_missing_deltakv_path=True only for construction-only tests." - ) - config.runtime_layout = RuntimeLayout.from_config(config.hf_config, require_mixed=is_qwen35) - if config.tiny_random: - if config.hf_config.num_attention_heads % config.tensor_parallel_size != 0: - raise ValueError( - "Tiny random num_attention_heads must be divisible by tensor_parallel_size." - ) - if config.hf_config.num_key_value_heads % config.tensor_parallel_size != 0: - raise ValueError( - "Tiny random num_key_value_heads must be divisible by tensor_parallel_size." - ) - if config.hf_config.vocab_size % config.tensor_parallel_size != 0: - raise ValueError( - "Tiny random vocab_size must be divisible by tensor_parallel_size." - ) if config.max_model_len > config.hf_config.max_position_embeddings: logger.warning('max_model_len > model.max_position_embeddings 输出可能不正常') config.hf_config.max_position_embeddings = config.max_model_len @@ -768,42 +75,36 @@ def _finalize_model_config(config, *, is_qwen35: bool) -> None: logger.warning('max_num_seqs_in_batch 过大或许会占用太多显存') -def load_and_validate_model(config) -> bool: +def load_and_validate_model(config) -> None: + validate_sparse_method_assets(config.vllm_sparse_method, config.model) if isinstance(config.deltakv_path, str): deltakv_path = config.deltakv_path.strip() - config.deltakv_path = None if deltakv_path.lower() in {"", "none", "null"} else deltakv_path + config.deltakv_path = ( + None + if deltakv_path.lower() in {"", "none", "null"} + else deltakv_path + ) if config.tiny_random and config.vllm_sparse_method == "deltakv": raise NotImplementedError( "Tiny random mode does not support DeltaKV compressor weights yet." ) - try: - config.outer_hf_config = AutoConfig.from_pretrained(config.model, trust_remote_code=True) - except Exception as e: - config.outer_hf_config = _load_raw_qwen35_config(config.model, e) - is_qwen35 = _is_qwen35_outer_config(config.outer_hf_config) + config.outer_hf_config = _load_model_config(config.model) + model_type = str(config_get(config.outer_hf_config, "model_type", "") or "") + model_spec = resolve_model_spec(model_type) config.hf_config = _extract_text_config(config.outer_hf_config) - if is_qwen35: - setattr(config.hf_config, "model_type", "qwen3_5") - model_type = str(getattr(config.hf_config, "model_type", "") or "") - is_minimax_m2 = model_type == "minimax_m2" - is_qwen3 = model_type == "qwen3" - is_qwen3_moe = model_type == "qwen3_moe" - - if config.vllm_sparse_method == "h2o": - if model_type not in H2O_SUPPORTED_MODEL_TYPES: - supported = ", ".join( - repr(value) for value in sorted(H2O_SUPPORTED_MODEL_TYPES) - ) - raise NotImplementedError( - "H2O v1 supports the model types already implemented by Sparse-vLLM: " - f"{supported}; got model_type={model_type!r}." - ) - + config.model_spec = model_spec + config.parallel_topology = model_spec.topology( + config.tensor_parallel_size, + config.expert_parallel_size, + config.data_parallel_size, + ) if config.tiny_random: from sparsevllm.debug.tiny_random import apply_tiny_random_overrides - if is_qwen35: - raise NotImplementedError("Tiny random mode does not support qwen3_5 yet.") + if not model_spec.supports_tiny_random: + raise NotImplementedError( + f"Tiny random mode does not support {model_spec.name} yet." + ) config.tiny_random_overrides = apply_tiny_random_overrides( config.hf_config, config.tiny_random_config, @@ -814,60 +115,46 @@ def load_and_validate_model(config) -> bool: f"seed={config.tiny_random_seed} overrides={config.tiny_random_overrides}", level="WARNING", ) + model_spec.validate_sharding(config.hf_config, config.parallel_topology) - raw_quantization_config = _config_get( + raw_quantization_config = config_get( config.hf_config, "quantization_config", - _config_get(config.outer_hf_config, "quantization_config", None), + config_get(config.outer_hf_config, "quantization_config", None), ) - quantized_model_name = "qwen3_5" - if is_minimax_m2: - quantized_model_name = "MiniMax M2.7" - elif is_qwen3: - quantized_model_name = "Qwen3" - elif is_qwen3_moe: - quantized_model_name = "Qwen3MoE" + use_unquantized_tiny = bool( + config.tiny_random and model_spec.tiny_random_unquantized + ) + if use_unquantized_tiny: + raw_quantization_config = None config.quantization_config = QuantizationConfig.from_hf_config( raw_quantization_config, - required_fp8=is_minimax_m2, - model_name=quantized_model_name, + required_fp8=model_spec.requires_fp8 and not use_unquantized_tiny, + model_name=model_spec.name, ) - if is_qwen35: - _validate_qwen35_checkpoint_precision( - config.hf_config, - raw_quantization_config, - config.quantization_config, - ) if config.tiny_random and config.quantization_config.enabled: - raise NotImplementedError("Tiny random mode does not support quantized model weights.") - setattr(config.hf_config, "quantization_config", config.quantization_config) - if is_minimax_m2: - _validate_minimax_m2_checkpoint_config( - config.hf_config, - raw_quantization_config, - ) - if is_qwen3 and config.quantization_config.enabled: - _validate_qwen3_fp8_checkpoint_config( - config.hf_config, - tensor_parallel_size=config.tensor_parallel_size, - ) - if is_qwen3_moe and config.quantization_config.enabled: - _validate_qwen3_moe_fp8_checkpoint_config( - config.hf_config, - raw_quantization_config, - ) - - if getattr(config.hf_config, "model_type", "") in {"deepseek_v2", "deepseek_v32"}: raise NotImplementedError( - f"Unsupported Sparse-vLLM model_type={config.hf_config.model_type!r}. " - "Supported model types: qwen2, qwen3, qwen3_5, llama." + "Tiny random mode does not support quantized model weights." + ) + setattr(config.hf_config, "quantization_config", config.quantization_config) + if not ( + config.tiny_random and model_spec.skip_checkpoint_validation_in_tiny_random + ): + validate_checkpoint( + model_type, + outer_config=config.outer_hf_config, + config=config.hf_config, + raw_quantization_config=raw_quantization_config, + quantization=config.quantization_config, + topology=config.parallel_topology, ) - if model_type == "qwen3_moe": - _validate_qwen3_moe_runtime(config, model_type=model_type) - elif model_type == "minimax_m2": - _validate_minimax_runtime(config, model_type=model_type) - else: - _validate_dense_parallelism(config, model_type=model_type) - _finalize_model_config(config, is_qwen35=is_qwen35) - return is_qwen35 + validate_model_runtime_compatibility( + model_type=model_type, + sparse_method=config.vllm_sparse_method, + topology=config.parallel_topology, + enforce_eager=config.enforce_eager, + decode_cuda_graph=config.decode_cuda_graph, + enable_prefix_caching=config.enable_prefix_caching, + ) + _finalize_model_config(config, model_spec) diff --git a/src/sparsevllm/configs/platform.py b/src/sparsevllm/configs/platform.py index 2777f3e6..7b21ffd1 100644 --- a/src/sparsevllm/configs/platform.py +++ b/src/sparsevllm/configs/platform.py @@ -2,8 +2,7 @@ import os -from sparsevllm.configs.common import _coerce_bool_config, _model_path_basename -from sparsevllm.configs.delta import SUPPORTED_SKIPKV_MODEL_NAMES +from sparsevllm.configs.common import _coerce_bool_config def _normalize_platform_aliases(config) -> None: @@ -39,14 +38,6 @@ def _normalize_platform_aliases(config) -> None: def normalize_platform(config) -> None: if not os.path.isdir(config.model): raise FileNotFoundError(f"Model directory does not exist: {config.model}") - if config.vllm_sparse_method == "skipkv": - model_name = _model_path_basename(config.model) - if model_name not in SUPPORTED_SKIPKV_MODEL_NAMES: - supported = ", ".join(sorted(SUPPORTED_SKIPKV_MODEL_NAMES)) - raise ValueError( - "SkipKV is supported only for the official models with released steering vectors: " - f"{supported}. Got model basename {model_name!r} from model path {config.model!r}." - ) config.tensor_parallel_size = int(config.tensor_parallel_size) config.expert_parallel_size = int(config.expert_parallel_size) config.data_parallel_size = int(config.data_parallel_size) diff --git a/src/sparsevllm/configs/prefix_cache.py b/src/sparsevllm/configs/prefix_cache.py index e3b488aa..db252397 100644 --- a/src/sparsevllm/configs/prefix_cache.py +++ b/src/sparsevllm/configs/prefix_cache.py @@ -1,4 +1,4 @@ -"""Prefix-cache normalization and model-dependent validation.""" +"""Prefix-cache normalization and validation.""" import math @@ -105,17 +105,22 @@ def normalize_prefix_cache(config) -> None: ) config.prefix_cache_salt = str(config.prefix_cache_salt or "") -def finalize_prefix_cache(config, *, is_qwen35: bool) -> None: +def finalize_prefix_cache(config) -> None: + block_multiple = config.model_spec.prefix_cache_block_size_multiple if ( - is_qwen35 + block_multiple is not None and config.resolved_prefix_cache_mode == "radix" and config.prefix_cache_block_size is None ): - config.prefix_cache_block_size = 4096 + config.prefix_cache_block_size = block_multiple config.prefix_cache_block_size = resolve_prefix_cache_block_size(config) - if is_qwen35 and config.resolved_prefix_cache_mode == "radix": - if config.prefix_cache_block_size < 4096 or config.prefix_cache_block_size % 4096 != 0: + if block_multiple is not None and config.resolved_prefix_cache_mode == "radix": + if ( + config.prefix_cache_block_size < block_multiple + or config.prefix_cache_block_size % block_multiple + ): raise ValueError( - "qwen3_5 mixed prefix cache requires prefix_cache_block_size to be " - f"4096*N, got {config.prefix_cache_block_size}." + f"{config.model_spec.name} prefix cache requires " + f"prefix_cache_block_size to be {block_multiple}*N, " + f"got {config.prefix_cache_block_size}." ) diff --git a/src/sparsevllm/configs/runtime.py b/src/sparsevllm/configs/runtime.py index f5ffd894..25459d59 100644 --- a/src/sparsevllm/configs/runtime.py +++ b/src/sparsevllm/configs/runtime.py @@ -1,7 +1,7 @@ """Top-level configuration composition and initialization orchestration.""" from dataclasses import dataclass, field -from typing import Any, Union +from typing import Any from sparsevllm.configs.bootstrap import normalize_bootstrap from sparsevllm.configs.cuda_graph import ( @@ -22,9 +22,6 @@ ) from sparsevllm.configs.model import ( AutoConfig, - Qwen3Config, - QuantizationConfig, - RuntimeLayout, load_and_validate_model, ) from sparsevllm.configs.platform import normalize_platform @@ -38,7 +35,11 @@ normalize_sparse_method_name, normalize_sparse_methods, ) +from sparsevllm.distributed import ParallelTopology from sparsevllm.method_registry import PREFILL_POLICY_AUTO +from sparsevllm.models.layout import RuntimeLayout +from sparsevllm.models.spec import ModelSpec +from sparsevllm.quantization import QuantizationConfig from sparsevllm.utils.log import logger @@ -70,10 +71,12 @@ class Config( # least one synchronous loading path when the budget is smaller. weight_loading_workers: int = 1 enforce_eager: bool = True - hf_config: Union[Qwen3Config, AutoConfig] | None = None + hf_config: AutoConfig | None = None outer_hf_config: Any | None = None runtime_layout: RuntimeLayout | None = None quantization_config: QuantizationConfig = field(default_factory=QuantizationConfig.disabled) + model_spec: ModelSpec = field(init=False, repr=False) + parallel_topology: ParallelTopology = field(init=False, repr=False) tiny_random: bool = False tiny_random_config: str | None = None tiny_random_seed: int = 0 @@ -84,14 +87,15 @@ class Config( @property def uses_outer_tp_moe_layout(self) -> bool: - model_type = str(getattr(self.hf_config, "model_type", "") or "") - return model_type in {"qwen3_moe", "minimax_m2"} and int( - self.tensor_parallel_size - ) > 1 + return self.parallel_topology.is_outer_tp_moe + + @property + def uses_dpa_ep_layout(self) -> bool: + return self.parallel_topology.is_dpa_ep @property def attention_tensor_parallel_size(self) -> int: - return int(self.tensor_parallel_size) + return self.parallel_topology.attention_tp_size @property def moe_expert_parallel_size(self) -> int: @@ -99,19 +103,11 @@ def moe_expert_parallel_size(self) -> int: @property def moe_tensor_parallel_size(self) -> int: - if self.uses_outer_tp_moe_layout: - return int(self.tensor_parallel_size) // int(self.expert_parallel_size) - return int(self.tensor_parallel_size) + return self.parallel_topology.moe_tp_size @property def world_size(self) -> int: - if self.uses_outer_tp_moe_layout: - return int(self.tensor_parallel_size) * int(self.data_parallel_size) - return ( - int(self.tensor_parallel_size) - * int(self.expert_parallel_size) - * int(self.data_parallel_size) - ) + return self.parallel_topology.world_size @property def weight_loading_workers_per_rank(self) -> int: @@ -128,10 +124,10 @@ def __post_init__(self): self, legacy_deltakv_graph_method=legacy_deltakv_graph_method, ) - is_qwen35 = load_and_validate_model(self) + load_and_validate_model(self) normalize_sparse_methods(self) - finalize_prefix_cache(self, is_qwen35=is_qwen35) - validate_deltakv_runtime(self, is_qwen35=is_qwen35) + finalize_prefix_cache(self) + validate_deltakv_runtime(self) finalize_sparse_layout(self) logger.info(f"LLM Config: {self}".replace('\n', ' ')) diff --git a/src/sparsevllm/configs/sparse.py b/src/sparsevllm/configs/sparse.py index 31286fc2..af47fcd2 100644 --- a/src/sparsevllm/configs/sparse.py +++ b/src/sparsevllm/configs/sparse.py @@ -5,8 +5,11 @@ _normalize_int_attr, _normalize_positive_int, ) -from sparsevllm.configs.delta import SUPPORTED_SKIPKV_MODEL_NAMES -from sparsevllm.method_registry import SUPPORTED_SPARSE_METHODS, normalize_sparse_method +from sparsevllm.method_registry import ( + SKIPKV_ASSET_MODEL_NAMES, + SUPPORTED_SPARSE_METHODS, + normalize_sparse_method, +) from sparsevllm.utils.log import logger, log_once def normalize_sparse_method_name(config) -> bool: @@ -168,7 +171,7 @@ def _normalize_skipkv(config) -> None: raise ValueError( "skipkv_enable_activation_steering=True requires skipkv_steering_vector_path. " "Official SkipKV support is limited to the released steering vectors for " - f"{', '.join(sorted(SUPPORTED_SKIPKV_MODEL_NAMES))}." + f"{', '.join(sorted(SKIPKV_ASSET_MODEL_NAMES))}." ) def normalize_sparse_methods(config) -> None: diff --git a/src/sparsevllm/debug/tiny_random.py b/src/sparsevllm/debug/tiny_random.py index 8a9fba23..27f80a3f 100644 --- a/src/sparsevllm/debug/tiny_random.py +++ b/src/sparsevllm/debug/tiny_random.py @@ -22,6 +22,17 @@ "num_hidden_layers", "num_key_value_heads", "vocab_size", + "q_lora_rank", + "o_lora_rank", + "o_groups", + "index_n_heads", + "index_head_dim", + "index_topk", + "num_local_experts", + "n_routed_experts", + "num_experts_per_tok", + "hc_mult", + "sliding_window", } ) @@ -106,21 +117,29 @@ def apply_tiny_random_overrides(hf_config: Any, path: str) -> dict[str, int]: f"{overrides[name]} > {original_values[name]}." ) - layer_types = getattr(hf_config, "layer_types", None) - if layer_types is not None: + for field_name in ("layer_types", "mlp_layer_types", "compress_ratios"): + layer_types = getattr(hf_config, field_name, None) + if layer_types is None: + continue num_layers = int(hf_config.num_hidden_layers) if len(layer_types) < num_layers: raise ValueError( - "Tiny random config cannot expand layer_types: " + f"Tiny random config cannot expand {field_name}: " f"requested={num_layers}, available={len(layer_types)}." ) - hf_config.layer_types = list(layer_types[:num_layers]) + setattr(hf_config, field_name, list(layer_types[:num_layers])) hidden_size = int(hf_config.hidden_size) num_heads = int(hf_config.num_attention_heads) num_kv_heads = int(hf_config.num_key_value_heads) head_dim = int(getattr(hf_config, "head_dim", hidden_size // num_heads)) - if hidden_size != num_heads * head_dim: + if getattr(hf_config, "model_type", "") == "deepseek_v4": + partial_rotary_factor = int(hf_config.qk_rope_head_dim) / head_dim + hf_config.partial_rotary_factor = partial_rotary_factor + for rope_parameters in getattr(hf_config, "rope_parameters", {}).values(): + if isinstance(rope_parameters, dict): + rope_parameters["partial_rotary_factor"] = partial_rotary_factor + if getattr(hf_config, "model_type", "") != "deepseek_v4" and hidden_size != num_heads * head_dim: raise ValueError( "Tiny random config requires hidden_size == num_attention_heads * head_dim, " f"got {hidden_size} != {num_heads} * {head_dim}." @@ -172,6 +191,7 @@ def initialize_sparse_model( from sparsevllm.utils.loader import ( _target_weight_name_for_model, default_weight_loader, + get_parameter_or_buffer, ) reference = build_tiny_random_hf_model(hf_config, seed=seed) @@ -204,7 +224,7 @@ def initialize_sparse_model( loaded_count += 1 break else: - param = model.get_parameter(param_name) + param = get_parameter_or_buffer(model, param_name) weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_parameter_names.add(param_name) diff --git a/src/sparsevllm/distributed/__init__.py b/src/sparsevllm/distributed/__init__.py index 9ce7f8c9..3f8700fc 100644 --- a/src/sparsevllm/distributed/__init__.py +++ b/src/sparsevllm/distributed/__init__.py @@ -2,22 +2,29 @@ ParallelContext, ParallelGroup, get_parallel_context, - hybrid_moe_group_ranks, init_parallel_context, + reset_parallel_context, +) +from sparsevllm.distributed.topology import ( + ParallelMode, + ParallelTopology, parallel_group_ranks, parallel_ranks_from_world_rank, - reset_parallel_context, world_rank_from_parallel_ranks, ) +from sparsevllm.distributed.sharding import validate_model_sharding, validate_top_k __all__ = [ "ParallelContext", "ParallelGroup", + "ParallelMode", + "ParallelTopology", "get_parallel_context", - "hybrid_moe_group_ranks", "init_parallel_context", "parallel_group_ranks", "parallel_ranks_from_world_rank", "reset_parallel_context", "world_rank_from_parallel_ranks", + "validate_model_sharding", + "validate_top_k", ] diff --git a/src/sparsevllm/distributed/parallel_context.py b/src/sparsevllm/distributed/parallel_context.py index 1304c8d5..cee2da05 100644 --- a/src/sparsevllm/distributed/parallel_context.py +++ b/src/sparsevllm/distributed/parallel_context.py @@ -1,152 +1,12 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field, replace import torch import torch.distributed as dist - -def _validate_sizes(tp_size: int, ep_size: int, dp_size: int) -> tuple[int, int, int]: - sizes = (int(tp_size), int(ep_size), int(dp_size)) - if any(size <= 0 for size in sizes): - raise ValueError( - "Parallel sizes must be positive, " - f"got TP={sizes[0]}, EP={sizes[1]}, DP={sizes[2]}." - ) - return sizes - - -def world_rank_from_parallel_ranks( - dp_rank: int, - ep_rank: int, - tp_rank: int, - *, - tp_size: int, - ep_size: int, - dp_size: int, -) -> int: - tp_size, ep_size, dp_size = _validate_sizes(tp_size, ep_size, dp_size) - dp_rank, ep_rank, tp_rank = int(dp_rank), int(ep_rank), int(tp_rank) - if not 0 <= dp_rank < dp_size: - raise ValueError(f"dp_rank must be in [0, {dp_size}), got {dp_rank}.") - if not 0 <= ep_rank < ep_size: - raise ValueError(f"ep_rank must be in [0, {ep_size}), got {ep_rank}.") - if not 0 <= tp_rank < tp_size: - raise ValueError(f"tp_rank must be in [0, {tp_size}), got {tp_rank}.") - return ((dp_rank * ep_size) + ep_rank) * tp_size + tp_rank - - -def parallel_ranks_from_world_rank( - world_rank: int, - *, - tp_size: int, - ep_size: int, - dp_size: int, -) -> tuple[int, int, int]: - tp_size, ep_size, dp_size = _validate_sizes(tp_size, ep_size, dp_size) - world_size = tp_size * ep_size * dp_size - world_rank = int(world_rank) - if not 0 <= world_rank < world_size: - raise ValueError(f"world_rank must be in [0, {world_size}), got {world_rank}.") - dp_ep_rank, tp_rank = divmod(world_rank, tp_size) - dp_rank, ep_rank = divmod(dp_ep_rank, ep_size) - return dp_rank, ep_rank, tp_rank - - -def parallel_group_ranks( - *, - tp_size: int, - ep_size: int, - dp_size: int, -) -> dict[str, tuple[tuple[int, ...], ...]]: - tp_size, ep_size, dp_size = _validate_sizes(tp_size, ep_size, dp_size) - - tensor_groups = tuple( - tuple( - world_rank_from_parallel_ranks( - dp_rank, - ep_rank, - tp_rank, - tp_size=tp_size, - ep_size=ep_size, - dp_size=dp_size, - ) - for tp_rank in range(tp_size) - ) - for dp_rank in range(dp_size) - for ep_rank in range(ep_size) - ) - expert_groups = tuple( - tuple( - world_rank_from_parallel_ranks( - dp_rank, - ep_rank, - tp_rank, - tp_size=tp_size, - ep_size=ep_size, - dp_size=dp_size, - ) - for ep_rank in range(ep_size) - ) - for dp_rank in range(dp_size) - for tp_rank in range(tp_size) - ) - data_groups = tuple( - tuple( - world_rank_from_parallel_ranks( - dp_rank, - ep_rank, - tp_rank, - tp_size=tp_size, - ep_size=ep_size, - dp_size=dp_size, - ) - for dp_rank in range(dp_size) - ) - for ep_rank in range(ep_size) - for tp_rank in range(tp_size) - ) - return { - "tensor": tensor_groups, - "expert": expert_groups, - "data": data_groups, - } - - -def hybrid_moe_group_ranks( - *, - outer_tp_size: int, - moe_ep_size: int, -) -> dict[str, tuple[tuple[int, ...], ...]]: - outer_tp_size = int(outer_tp_size) - moe_ep_size = int(moe_ep_size) - if outer_tp_size <= 0 or moe_ep_size <= 0: - raise ValueError( - "Hybrid MoE parallel sizes must be positive, " - f"got outer TP={outer_tp_size}, MoE EP={moe_ep_size}." - ) - if outer_tp_size % moe_ep_size: - raise ValueError( - "Hybrid MoE outer TP must be divisible by MoE EP, " - f"got outer TP={outer_tp_size}, MoE EP={moe_ep_size}." - ) - moe_tp_size = outer_tp_size // moe_ep_size - attention_groups = (tuple(range(outer_tp_size)),) - moe_tensor_groups = tuple( - tuple(range(ep_rank * moe_tp_size, (ep_rank + 1) * moe_tp_size)) - for ep_rank in range(moe_ep_size) - ) - moe_expert_groups = tuple( - tuple(ep_rank * moe_tp_size + moe_tp_rank for ep_rank in range(moe_ep_size)) - for moe_tp_rank in range(moe_tp_size) - ) - singleton_groups = tuple((rank,) for rank in range(outer_tp_size)) - return { - "attention": attention_groups, - "moe_tensor": moe_tensor_groups, - "moe_expert": moe_expert_groups, - "data": singleton_groups, - } +from sparsevllm.distributed.topology import ParallelTopology, parallel_group_ranks +from sparsevllm.operators.all_reduce import AllReduceProvider, resolve_all_reduce_provider @dataclass(frozen=True) @@ -155,6 +15,7 @@ class ParallelGroup: ranks: tuple[int, ...] rank: int size: int + all_reduce_provider: AllReduceProvider | None = field(default=None, compare=False, repr=False) def __post_init__(self) -> None: if self.size != len(self.ranks): @@ -234,7 +95,10 @@ def _all_reduce( op: dist.ReduceOp = dist.ReduceOp.SUM, ) -> torch.Tensor: if group.size > 1: - dist.all_reduce(tensor, op=op, group=group.process_group) + if op != dist.ReduceOp.SUM or group.all_reduce_provider is None: + dist.all_reduce(tensor, op=op, group=group.process_group) + else: + tensor = group.all_reduce_provider.run(tensor) return tensor def world_all_reduce( @@ -265,6 +129,41 @@ def ep_all_reduce( ) -> torch.Tensor: return self._all_reduce(tensor, self.expert, op) + def ep_all_gather_into_tensor(self, tensor: torch.Tensor) -> torch.Tensor: + if self.ep_size == 1: + return tensor + output = torch.empty( + (self.ep_size * tensor.shape[0], *tensor.shape[1:]), + dtype=tensor.dtype, + device=tensor.device, + ) + dist.all_gather_into_tensor( + output, + tensor.contiguous(), + group=self.expert.process_group, + ) + return output + + def ep_reduce_scatter_tensor(self, tensor: torch.Tensor) -> torch.Tensor: + if self.ep_size == 1: + return tensor + if tensor.shape[0] % self.ep_size: + raise ValueError( + "EP reduce-scatter requires dim 0 divisible by EP size: " + f"{tensor.shape[0]} % {self.ep_size}." + ) + output = torch.empty( + (tensor.shape[0] // self.ep_size, *tensor.shape[1:]), + dtype=tensor.dtype, + device=tensor.device, + ) + dist.reduce_scatter_tensor( + output, + tensor.contiguous(), + group=self.expert.process_group, + ) + return output + def moe_tp_all_reduce( self, tensor: torch.Tensor, @@ -339,10 +238,7 @@ def _local_group( def init_parallel_context( *, - tp_size: int, - ep_size: int, - dp_size: int, - hybrid_moe: bool = False, + topology: ParallelTopology, ) -> ParallelContext: global _PARALLEL_CONTEXT if _PARALLEL_CONTEXT is not None: @@ -350,10 +246,10 @@ def init_parallel_context( if not dist.is_initialized(): raise RuntimeError("torch.distributed must be initialized before ParallelContext.") - tp_size, ep_size, dp_size = _validate_sizes(tp_size, ep_size, dp_size) - if hybrid_moe and dp_size != 1: - raise ValueError(f"Hybrid MoE parallelism requires DP=1, got DP={dp_size}.") - expected_world_size = tp_size if hybrid_moe else tp_size * ep_size * dp_size + tp_size = topology.tensor_parallel_size + ep_size = topology.expert_parallel_size + dp_size = topology.data_parallel_size + expected_world_size = topology.world_size world_size = dist.get_world_size() world_rank = dist.get_rank() if world_size != expected_world_size: @@ -362,24 +258,7 @@ def init_parallel_context( f"world_size={world_size}, TP={tp_size}, EP={ep_size}, DP={dp_size}." ) - if hybrid_moe: - hybrid_groups = hybrid_moe_group_ranks( - outer_tp_size=tp_size, - moe_ep_size=ep_size, - ) - ranks_by_dimension = { - "tensor": hybrid_groups["attention"], - "expert": hybrid_groups["moe_expert"], - "data": hybrid_groups["data"], - "moe_tensor": hybrid_groups["moe_tensor"], - } - else: - ranks_by_dimension = parallel_group_ranks( - tp_size=tp_size, - ep_size=ep_size, - dp_size=dp_size, - ) - ranks_by_dimension["moe_tensor"] = ranks_by_dimension["tensor"] + ranks_by_dimension = parallel_group_ranks(topology) world_ranks = tuple(range(world_size)) process_groups: dict[tuple[int, ...], dist.ProcessGroup | None] = { world_ranks: dist.group.WORLD, @@ -390,7 +269,7 @@ def init_parallel_context( continue process_groups[ranks] = None if len(ranks) == 1 else dist.new_group(list(ranks)) - _PARALLEL_CONTEXT = ParallelContext( + context = ParallelContext( world=ParallelGroup( process_group=dist.group.WORLD, ranks=world_ranks, @@ -404,6 +283,25 @@ def init_parallel_context( ranks_by_dimension["moe_tensor"], process_groups, world_rank ), ) + providers: dict[tuple[int, ...], AllReduceProvider] = {} + + def bind_provider(group: ParallelGroup) -> ParallelGroup: + if group.size == 1: + return group + provider = providers.get(group.ranks) + if provider is None: + provider = providers[group.ranks] = resolve_all_reduce_provider( + group.process_group, group.size + ) + return replace(group, all_reduce_provider=provider) + + _PARALLEL_CONTEXT = ParallelContext( + world=bind_provider(context.world), + tensor=bind_provider(context.tensor), + expert=bind_provider(context.expert), + data=bind_provider(context.data), + moe_tensor=bind_provider(context.moe_tensor or context.tensor), + ) return _PARALLEL_CONTEXT diff --git a/src/sparsevllm/distributed/sharding.py b/src/sparsevllm/distributed/sharding.py new file mode 100644 index 00000000..bdb2ff41 --- /dev/null +++ b/src/sparsevllm/distributed/sharding.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from collections.abc import Mapping + +from sparsevllm.distributed.topology import ParallelTopology + + +def _validate_divisible( + model_name: str, + fields: Mapping[str, int], + divisor: int, + parallelism: str, +) -> None: + invalid = { + name: int(value) + for name, value in fields.items() + if int(value) <= 0 or int(value) % divisor + } + if invalid: + raise ValueError( + f"{model_name} dimensions must be positive and divisible by " + f"{parallelism}={divisor}, invalid={invalid}." + ) + + +def validate_model_sharding( + topology: ParallelTopology, + *, + model_name: str, + attention_fields: Mapping[str, int], + num_experts: int | None = None, + moe_fields: Mapping[str, int] | None = None, +) -> None: + _validate_divisible( + model_name, + attention_fields, + topology.attention_tp_size, + "attention TP", + ) + if num_experts is None: + return + num_experts = int(num_experts) + if num_experts <= 0 or num_experts % topology.expert_parallel_size: + raise ValueError( + f"{model_name} num_experts must be positive and divisible by " + f"EP={topology.expert_parallel_size}, got {num_experts}." + ) + if moe_fields: + _validate_divisible( + model_name, + moe_fields, + topology.moe_tp_size, + "MoE TP", + ) + + +def validate_top_k(model_name: str, top_k: int, num_experts: int) -> None: + if not 1 <= int(top_k) <= int(num_experts): + raise ValueError( + f"{model_name} top_k must be in [1, num_experts], " + f"got top_k={top_k}, num_experts={num_experts}." + ) diff --git a/src/sparsevllm/distributed/topology.py b/src/sparsevllm/distributed/topology.py new file mode 100644 index 00000000..327f6e9e --- /dev/null +++ b/src/sparsevllm/distributed/topology.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + + +class ParallelMode(str, Enum): + STANDARD = "standard" + OUTER_TP_MOE = "outer_tp_moe_tp_ep" + DPA_EP = "dpa_ep" + + +@dataclass(frozen=True) +class ParallelTopology: + tensor_parallel_size: int + expert_parallel_size: int + data_parallel_size: int + mode: ParallelMode = ParallelMode.STANDARD + + def __post_init__(self) -> None: + for field in ( + "tensor_parallel_size", + "expert_parallel_size", + "data_parallel_size", + ): + object.__setattr__(self, field, int(getattr(self, field))) + object.__setattr__(self, "mode", ParallelMode(self.mode)) + sizes = ( + self.tensor_parallel_size, + self.expert_parallel_size, + self.data_parallel_size, + ) + if any(size <= 0 for size in sizes): + raise ValueError( + "Parallel sizes must be positive, " + f"got TP={sizes[0]}, EP={sizes[1]}, DP={sizes[2]}." + ) + if self.mode is ParallelMode.OUTER_TP_MOE: + if self.data_parallel_size != 1: + raise ValueError( + "Outer-TP MoE parallelism requires DP=1, " + f"got DP={self.data_parallel_size}." + ) + if self.tensor_parallel_size % self.expert_parallel_size: + raise ValueError( + "Outer-TP MoE requires TP divisible by EP, " + f"got TP={self.tensor_parallel_size}, EP={self.expert_parallel_size}." + ) + elif self.mode is ParallelMode.DPA_EP and ( + self.tensor_parallel_size != 1 + or self.expert_parallel_size != self.data_parallel_size + ): + raise ValueError( + "DPA+EP requires TP=1 and matching DP/EP sizes, " + f"got TP={self.tensor_parallel_size}, " + f"EP={self.expert_parallel_size}, DP={self.data_parallel_size}." + ) + + @property + def is_outer_tp_moe(self) -> bool: + return self.mode is ParallelMode.OUTER_TP_MOE + + @property + def is_dpa_ep(self) -> bool: + return self.mode is ParallelMode.DPA_EP + + @property + def attention_tp_size(self) -> int: + return 1 if self.is_dpa_ep else self.tensor_parallel_size + + @property + def moe_tp_size(self) -> int: + if self.is_dpa_ep: + return 1 + return ( + self.tensor_parallel_size // self.expert_parallel_size + if self.is_outer_tp_moe + else self.tensor_parallel_size + ) + + @property + def world_size(self) -> int: + if self.is_dpa_ep: + return self.data_parallel_size + return ( + self.tensor_parallel_size + if self.is_outer_tp_moe + else self.tensor_parallel_size + * self.expert_parallel_size + * self.data_parallel_size + ) + + +def world_rank_from_parallel_ranks( + topology: ParallelTopology, + dp_rank: int, + ep_rank: int, + tp_rank: int, +) -> int: + if topology.mode is not ParallelMode.STANDARD: + raise ValueError( + f"{topology.mode.value} does not use standard DP/EP/TP rank mapping." + ) + dp_rank, ep_rank, tp_rank = int(dp_rank), int(ep_rank), int(tp_rank) + for name, rank, size in ( + ("dp_rank", dp_rank, topology.data_parallel_size), + ("ep_rank", ep_rank, topology.expert_parallel_size), + ("tp_rank", tp_rank, topology.tensor_parallel_size), + ): + if not 0 <= rank < size: + raise ValueError(f"{name} must be in [0, {size}), got {rank}.") + return ( + (dp_rank * topology.expert_parallel_size + ep_rank) + * topology.tensor_parallel_size + + tp_rank + ) + + +def parallel_ranks_from_world_rank( + topology: ParallelTopology, + world_rank: int, +) -> tuple[int, int, int]: + if topology.mode is not ParallelMode.STANDARD: + raise ValueError( + f"{topology.mode.value} does not use standard DP/EP/TP rank mapping." + ) + world_rank = int(world_rank) + if not 0 <= world_rank < topology.world_size: + raise ValueError( + f"world_rank must be in [0, {topology.world_size}), got {world_rank}." + ) + dp_ep_rank, tp_rank = divmod(world_rank, topology.tensor_parallel_size) + dp_rank, ep_rank = divmod(dp_ep_rank, topology.expert_parallel_size) + return dp_rank, ep_rank, tp_rank + + +def _standard_group_ranks( + topology: ParallelTopology, +) -> dict[str, tuple[tuple[int, ...], ...]]: + tp_size = topology.tensor_parallel_size + ep_size = topology.expert_parallel_size + dp_size = topology.data_parallel_size + + def world_rank(dp_rank: int, ep_rank: int, tp_rank: int) -> int: + return world_rank_from_parallel_ranks(topology, dp_rank, ep_rank, tp_rank) + + tensor_groups = tuple( + tuple(world_rank(dp_rank, ep_rank, tp_rank) for tp_rank in range(tp_size)) + for dp_rank in range(dp_size) + for ep_rank in range(ep_size) + ) + return { + "tensor": tensor_groups, + "expert": tuple( + tuple(world_rank(dp_rank, ep_rank, tp_rank) for ep_rank in range(ep_size)) + for dp_rank in range(dp_size) + for tp_rank in range(tp_size) + ), + "data": tuple( + tuple(world_rank(dp_rank, ep_rank, tp_rank) for dp_rank in range(dp_size)) + for ep_rank in range(ep_size) + for tp_rank in range(tp_size) + ), + "moe_tensor": tensor_groups, + } + + +def _outer_tp_moe_group_ranks( + topology: ParallelTopology, +) -> dict[str, tuple[tuple[int, ...], ...]]: + outer_tp_size = topology.attention_tp_size + moe_ep_size = topology.expert_parallel_size + moe_tp_size = topology.moe_tp_size + return { + "tensor": (tuple(range(outer_tp_size)),), + "expert": tuple( + tuple( + ep_rank * moe_tp_size + moe_tp_rank + for ep_rank in range(moe_ep_size) + ) + for moe_tp_rank in range(moe_tp_size) + ), + "data": tuple((rank,) for rank in range(outer_tp_size)), + "moe_tensor": tuple( + tuple(range(ep_rank * moe_tp_size, (ep_rank + 1) * moe_tp_size)) + for ep_rank in range(moe_ep_size) + ), + } + + +def _dpa_ep_group_ranks( + topology: ParallelTopology, +) -> dict[str, tuple[tuple[int, ...], ...]]: + world = tuple(range(topology.world_size)) + singletons = tuple((rank,) for rank in world) + return { + "tensor": singletons, + "expert": (world,), + "data": (world,), + "moe_tensor": singletons, + } + + +def parallel_group_ranks( + topology: ParallelTopology, +) -> dict[str, tuple[tuple[int, ...], ...]]: + if topology.is_dpa_ep: + return _dpa_ep_group_ranks(topology) + if topology.is_outer_tp_moe: + return _outer_tp_moe_group_ranks(topology) + return _standard_group_ranks(topology) diff --git a/src/sparsevllm/engine/cache_manager/__init__.py b/src/sparsevllm/engine/cache_manager/__init__.py index 79a1a02b..68ab28bf 100644 --- a/src/sparsevllm/engine/cache_manager/__init__.py +++ b/src/sparsevllm/engine/cache_manager/__init__.py @@ -20,6 +20,7 @@ "DeltaKVCacheTritonManagerV4", "DeltaKVLessMemoryCacheManager", "DeltaKVLessMemoryCudaGraphCacheManager", + "DeepseekV4CacheManager", ] @@ -72,5 +73,9 @@ def __getattr__(name: str): from .deltakv_less_memory_cuda_graph import DeltaKVLessMemoryCudaGraphCacheManager return DeltaKVLessMemoryCudaGraphCacheManager + if name == "DeepseekV4CacheManager": + from .deepseek_v4 import DeepseekV4CacheManager + + return DeepseekV4CacheManager raise AttributeError(name) diff --git a/src/sparsevllm/engine/cache_manager/base.py b/src/sparsevllm/engine/cache_manager/base.py index 035eb315..cc8bc6a7 100644 --- a/src/sparsevllm/engine/cache_manager/base.py +++ b/src/sparsevllm/engine/cache_manager/base.py @@ -191,9 +191,7 @@ def __init__(self, config: Config, parallel_context: ParallelContext): self.num_layers = self.hf_config.num_hidden_layers self.runtime_layout = getattr(config, "runtime_layout", None) if self.runtime_layout is None: - from sparsevllm.config import RuntimeLayout - - self.runtime_layout = RuntimeLayout.dense(self.num_layers) + raise ValueError("CacheManager requires config.runtime_layout.") self.num_kv_layers = int(self.runtime_layout.num_kv_layers) self.num_kv_heads = self.hf_config.num_key_value_heads // self.tp_size @@ -268,12 +266,24 @@ def _is_stream_capturing(self) -> bool: @staticmethod def create(config: Config, parallel_context: ParallelContext) -> "CacheManager": sparse_method = normalize_sparse_method(config.vllm_sparse_method) - model_type = getattr(getattr(config, "hf_config", None), "model_type", "") or "" - - if model_type in {"deepseek_v2", "deepseek_v32"}: - raise NotImplementedError(f"Unsupported Sparse-vLLM model_type={model_type!r}.") if sparse_method not in SUPPORTED_SPARSE_METHODS: raise ValueError(f"Unsupported vllm_sparse_method={sparse_method!r}.") + cache_manager_class_name = str( + getattr( + getattr(config, "model_spec", None), + "cache_manager_class_name", + "", + ) + or "" + ) + if cache_manager_class_name: + from importlib import import_module + + cache_manager_class = getattr( + import_module("sparsevllm.engine.cache_manager"), + cache_manager_class_name, + ) + return cache_manager_class(config, parallel_context) if sparse_method == "deltakv": from .deltakv_runtime import DeltaKVCacheManager diff --git a/src/sparsevllm/engine/cache_manager/deepseek_v4.py b/src/sparsevllm/engine/cache_manager/deepseek_v4.py new file mode 100644 index 00000000..8016a6d2 --- /dev/null +++ b/src/sparsevllm/engine/cache_manager/deepseek_v4.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import math + +import torch + +from sparsevllm.utils.log import logger + +from .standard import StandardCacheManager + + +class DeepseekV4CacheManager(StandardCacheManager): + """Row-indexed DSV4 sliding and compressed KV storage. + + DSV4 does not use a conventional K/V pair per source token. Every layer + keeps a 128-token shared-KV ring, while CSA and HCA layers additionally + retain one compressed entry per 4 or 128 source tokens. Keeping these + tensors here makes request ownership and CUDA-graph metadata follow the + same cache-manager contract as the other runtimes. + """ + + def allocate_kv_cache(self) -> None: + config = self.config + hf_config = self.hf_config + rows = int(self.max_buffer_rows) + max_len = int(self.max_model_len) + window = int(hf_config.sliding_window) + head_dim = int(hf_config.head_dim) + index_dim = int(hf_config.index_head_dim) + csa_ratio = int(hf_config.compress_rates["compressed_sparse_attention"]) + hca_ratio = int(hf_config.compress_rates["heavily_compressed_attention"]) + layer_types = tuple(hf_config.layer_types) + self.csa_layers = tuple(i for i, kind in enumerate(layer_types) if kind == "compressed_sparse_attention") + self.hca_layers = tuple(i for i, kind in enumerate(layer_types) if kind == "heavily_compressed_attention") + self._csa_slot = {layer: i for i, layer in enumerate(self.csa_layers)} + self._hca_slot = {layer: i for i, layer in enumerate(self.hca_layers)} + self.csa_ratio = csa_ratio + self.hca_ratio = hca_ratio + self.sliding_window = window + self.max_csa_entries = math.ceil(max_len / csa_ratio) + self.max_hca_entries = math.ceil(max_len / hca_ratio) + dtype = hf_config.torch_dtype + + shapes = { + "raw_kv": (len(layer_types), rows, window, head_dim), + "csa_kv": (len(self.csa_layers), rows, self.max_csa_entries, head_dim), + "csa_index": (len(self.csa_layers), rows, self.max_csa_entries, index_dim), + "hca_kv": (len(self.hca_layers), rows, self.max_hca_entries, head_dim), + "csa_ring_kv": (len(self.csa_layers), rows, csa_ratio, 2 * head_dim), + "csa_ring_gate": (len(self.csa_layers), rows, csa_ratio, 2 * head_dim), + "csa_overlap_kv": (len(self.csa_layers), rows, csa_ratio, head_dim), + "csa_overlap_gate": (len(self.csa_layers), rows, csa_ratio, head_dim), + "index_ring_kv": (len(self.csa_layers), rows, csa_ratio, 2 * index_dim), + "index_ring_gate": (len(self.csa_layers), rows, csa_ratio, 2 * index_dim), + "index_overlap_kv": (len(self.csa_layers), rows, csa_ratio, index_dim), + "index_overlap_gate": (len(self.csa_layers), rows, csa_ratio, index_dim), + "hca_ring_kv": (len(self.hca_layers), rows, hca_ratio, head_dim), + "hca_ring_gate": (len(self.hca_layers), rows, hca_ratio, head_dim), + } + element_size = torch.empty((), dtype=dtype).element_size() + tensor_bytes = sum(math.prod(shape) for shape in shapes.values()) * element_size + metadata_bytes = 2 * rows * max_len * torch.empty((), dtype=torch.int32).element_size() + cache_bytes = tensor_bytes + metadata_bytes + free, total = self.platform.get_available_memory(self.device.index or 0) + budget = max(0, int(total * float(config.gpu_memory_utilization)) - int(total - free)) + if cache_bytes > budget: + raise RuntimeError( + "DeepSeek V4 cache does not fit the configured GPU budget: " + f"required={cache_bytes / 2**30:.2f} GiB available={budget / 2**30:.2f} GiB, " + f"max_model_len={max_len}, max_num_seqs_in_gpu={rows}. Reduce one of these limits." + ) + + def empty(name: str) -> torch.Tensor: + return torch.empty(shapes[name], dtype=dtype, device=self.device) + + self.raw_kv = empty("raw_kv") + self.csa_kv = empty("csa_kv") + self.csa_index = empty("csa_index") + self.hca_kv = empty("hca_kv") + self.csa_ring_kv = empty("csa_ring_kv") + self.csa_ring_gate = empty("csa_ring_gate") + self.csa_overlap_kv = empty("csa_overlap_kv") + self.csa_overlap_gate = empty("csa_overlap_gate") + self.index_ring_kv = empty("index_ring_kv") + self.index_ring_gate = empty("index_ring_gate") + self.index_overlap_kv = empty("index_overlap_kv") + self.index_overlap_gate = empty("index_overlap_gate") + self.hca_ring_kv = empty("hca_ring_kv") + self.hca_ring_gate = empty("hca_ring_gate") + self.reset_deepseek_v4_cache() + + # StandardCacheManager still owns admission and request-to-row metadata. + # Its token slots are bookkeeping-only for DSV4; attention reads the + # row-indexed tensors above. + config.num_kvcache_slots = rows * max_len + self.kv_cache = torch.empty(0, dtype=dtype, device=self.device) + logger.info( + "DeepSeek V4 cache: {:.2f} GiB, rows={}, max_len={}, CSA={}, HCA={}.", + cache_bytes / 2**30, + rows, + max_len, + len(self.csa_layers), + len(self.hca_layers), + ) + + def reset_deepseek_v4_cache(self) -> None: + for tensor in ( + self.raw_kv, + self.csa_kv, + self.csa_index, + self.hca_kv, + self.csa_ring_kv, + self.csa_overlap_kv, + self.index_ring_kv, + self.index_overlap_kv, + self.hca_ring_kv, + ): + tensor.zero_() + for tensor in ( + self.csa_ring_gate, + self.csa_overlap_gate, + self.index_ring_gate, + self.index_overlap_gate, + self.hca_ring_gate, + ): + tensor.fill_(float("-inf")) + + def reset_after_warmup(self) -> None: + super().reset_after_warmup() + self.reset_deepseek_v4_cache() + + def free_seq(self, seq_id: int) -> None: + row = self.seq_id_to_row.get(int(seq_id)) + super().free_seq(seq_id) + if row is None: + return + self.csa_overlap_kv[:, row].zero_() + self.csa_overlap_gate[:, row].fill_(float("-inf")) + self.index_overlap_kv[:, row].zero_() + self.index_overlap_gate[:, row].fill_(float("-inf")) + + def csa_slot(self, layer_idx: int) -> int: + return self._csa_slot[int(layer_idx)] + + def hca_slot(self, layer_idx: int) -> int: + return self._hca_slot[int(layer_idx)] + + def compressed_capacity(self, ratio: int, positions: torch.Tensor) -> int: + static_len = self._decode_static_max_context_len + if static_len is not None: + return min(math.ceil(int(static_len) / int(ratio)), math.ceil(self.max_model_len / int(ratio))) + return min(math.ceil((int(positions.max()) + 1) / int(ratio)), math.ceil(self.max_model_len / int(ratio))) + + def decode_cuda_graph_keepalive_tensors(self) -> list[torch.Tensor]: + return [ + self.raw_kv, + self.csa_kv, + self.csa_index, + self.hca_kv, + self.csa_ring_kv, + self.csa_ring_gate, + self.csa_overlap_kv, + self.csa_overlap_gate, + self.index_ring_kv, + self.index_ring_gate, + self.index_overlap_kv, + self.index_overlap_gate, + self.hca_ring_kv, + self.hca_ring_gate, + ] diff --git a/src/sparsevllm/engine/input_processor.py b/src/sparsevllm/engine/input_processor.py new file mode 100644 index 00000000..5a6ccee7 --- /dev/null +++ b/src/sparsevllm/engine/input_processor.py @@ -0,0 +1,29 @@ +from transformers import PreTrainedTokenizerBase + + +def tokenize_text_prompt( + tokenizer: PreTrainedTokenizerBase | None, + prompt: str | list[int], +) -> list[int]: + """Validate and tokenize the text-only engine input contract.""" + if isinstance(prompt, str): + if tokenizer is None: + raise RuntimeError("A tokenizer is required for string prompts.") + add_special_tokens = ( + tokenizer.bos_token is not None + and not prompt.startswith(tokenizer.bos_token) + ) + return tokenizer.encode(prompt, add_special_tokens=add_special_tokens) + + if not isinstance(prompt, list) or not all( + isinstance(token_id, int) and not isinstance(token_id, bool) + for token_id in prompt + ): + raise TypeError( + "Sparse-vLLM accepts text prompts only: prompt must be a string " + "or a flat list of integer token IDs; structured image, video, " + "or MTP request objects are unsupported." + ) + if not prompt: + raise ValueError("prompt token IDs must not be empty.") + return list(prompt) diff --git a/src/sparsevllm/engine/llm_engine.py b/src/sparsevllm/engine/llm_engine.py index bcb46146..a18b3465 100644 --- a/src/sparsevllm/engine/llm_engine.py +++ b/src/sparsevllm/engine/llm_engine.py @@ -19,6 +19,7 @@ from sparsevllm.engine.sequence import Sequence from sparsevllm.engine.scheduler import Scheduler from sparsevllm.engine.model_runner import ModelRunner, make_tp_shm_name +from sparsevllm.engine.input_processor import tokenize_text_prompt from sparsevllm.engine.prefix_cache import PrefixCacheRoutingSnapshot from sparsevllm.engine.chain_cache import ( ChainCacheIndex, @@ -61,8 +62,7 @@ def _use_graph_scaled_warmup(config: Config) -> bool: def _moe_workspace_warmup_token_counts(config: Config) -> tuple[int, ...]: - model_type = str(getattr(config.hf_config, "model_type", "") or "") - if model_type not in {"qwen3_moe", "minimax_m2"}: + if config.model_spec.num_experts_field is None: return () max_batched_tokens = int(config.max_num_batched_tokens) @@ -453,6 +453,7 @@ def run_warmup( self._warmup_moe_workspaces() self._after_warmup_debug_cleanup() + self.model_runner.call("log_operator_implementations") logger.info("Warmup finished.") def _warmup_moe_workspaces(self) -> None: @@ -574,16 +575,8 @@ def call_model_runner_exit(): return runner_exit_completed, runner_platform def _tokenize_prompt(self, prompt: str | list[int]) -> list[int]: - if isinstance(prompt, str): - # Add BOS for raw prompts, but do not duplicate it when a chat - # template already starts with BOS. - add_special_tokens = True - if self.tokenizer.bos_token is None or prompt.startswith(self.tokenizer.bos_token): - add_special_tokens = False - return self.tokenizer.encode( - prompt, add_special_tokens=add_special_tokens - ) - return [int(token_id) for token_id in prompt] + tokenizer = self.tokenizer if isinstance(prompt, str) else None + return tokenize_text_prompt(tokenizer, prompt) def admit_request( self, diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 7f311589..bd0eb0b1 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -18,6 +18,7 @@ from sparsevllm.models.qwen2 import Qwen2ForCausalLM from sparsevllm.models.llama import LlamaForCausalLM from sparsevllm.layers.sampler import Sampler +from sparsevllm.operators import registry as operator_registry from sparsevllm.utils.context import set_context, get_context, reset_context from sparsevllm.utils.loader import load_model, sync_deltakv_config_from_checkpoint @@ -29,6 +30,7 @@ from sparsevllm.engine.recurrent_state_manager import RecurrentStateManager, RecurrentStateSpec from sparsevllm.engine.runtime_state import RuntimeState from sparsevllm.engine.sparse_controller import SparseController +from sparsevllm.models.spec import ModelSpec import sparsevllm.platforms as platforms from sparsevllm.utils.profiler import profiler @@ -44,17 +46,31 @@ try: from sparsevllm.models.minimax_m2 import MiniMaxM2ForCausalLM - _MINIMAX_M2_IMPORT_ERROR = None -except ImportError as exc: +except ImportError: MiniMaxM2ForCausalLM = None - _MINIMAX_M2_IMPORT_ERROR = exc try: from sparsevllm.models.qwen3_5 import Qwen35ForCausalLM - _QWEN35_IMPORT_ERROR = None -except ImportError as exc: +except ImportError: Qwen35ForCausalLM = None - _QWEN35_IMPORT_ERROR = exc + +try: + from sparsevllm.models.qwen3_5_moe import Qwen35MoeForCausalLM +except ImportError: + Qwen35MoeForCausalLM = None + +try: + from sparsevllm.models.deepseek_v4 import DeepseekV4ForCausalLM +except ImportError: + DeepseekV4ForCausalLM = None + + +def _create_model(hf_config, model_spec: ModelSpec): + class_name = model_spec.runtime_class_name + model_class = globals().get(class_name) + if model_class is None: + raise ImportError(f"{class_name} is unavailable for {model_spec.name}.") + return model_class(hf_config) TP_SHM_NAME_PREFIX = "sparsevllm_" @@ -78,6 +94,7 @@ "debug_moe_states_cpu", "free_slots", "free_slots_batch", + "log_operator_implementations", "refresh_prefix_cache_hit", "reset_after_warmup", "run", @@ -131,10 +148,7 @@ def __init__( rank=rank, ) self.parallel_context = init_parallel_context( - tp_size=config.tensor_parallel_size, - ep_size=config.expert_parallel_size, - dp_size=config.data_parallel_size, - hybrid_moe=config.uses_outer_tp_moe_layout, + topology=config.parallel_topology, ) # CUDA allocator peaks are process-global and survive LLMEngine.exit(). @@ -146,43 +160,18 @@ def __init__( torch.set_default_dtype(hf_config.torch_dtype) torch.set_default_device(self.device) setattr(hf_config, "mlp_chunk_size", config.mlp_chunk_size) + setattr( + hf_config, + "decode_cuda_graph", + bool(getattr(config, "decode_cuda_graph", False)), + ) + setattr( + hf_config, + "sparsevllm_tiny_random", + bool(getattr(config, "tiny_random", False)), + ) - # 加载对应的模型分片 (Shards) - if hf_config.model_type == "qwen2": - self.model = Qwen2ForCausalLM(hf_config) - elif hf_config.model_type == "qwen3": - if Qwen3ForCausalLM is None: - raise ImportError( - "Qwen3ForCausalLM is unavailable in this Transformers installation. " - "Use a Transformers version with Qwen3 support for Qwen3 models." - ) - self.model = Qwen3ForCausalLM(hf_config) - elif hf_config.model_type == "qwen3_moe": - if Qwen3MoeForCausalLM is None: - raise ImportError( - "Qwen3MoeForCausalLM is unavailable in this Transformers installation. " - "Use a Transformers version with Qwen3MoE config support." - ) - self.model = Qwen3MoeForCausalLM(hf_config) - elif hf_config.model_type == "minimax_m2": - if MiniMaxM2ForCausalLM is None: - raise ImportError( - "MiniMaxM2ForCausalLM is unavailable; verify the MiniMax FP8 " - "runtime dependencies in the active uv environment: " - f"{_MINIMAX_M2_IMPORT_ERROR}" - ) from _MINIMAX_M2_IMPORT_ERROR - self.model = MiniMaxM2ForCausalLM(hf_config) - elif hf_config.model_type == "qwen3_5": - if Qwen35ForCausalLM is None: - raise ImportError( - "Qwen35ForCausalLM is unavailable. Install the qwen3_5 runtime " - f"dependencies and verify vendored kernels import correctly: {_QWEN35_IMPORT_ERROR}" - ) from _QWEN35_IMPORT_ERROR - self.model = Qwen35ForCausalLM(hf_config) - elif hf_config.model_type == "llama": - self.model = LlamaForCausalLM(hf_config) - else: - raise NotImplementedError(f"Unsupported Sparse-vLLM model_type={hf_config.model_type!r}.") + self.model = _create_model(hf_config, config.model_spec) if config.tiny_random: from sparsevllm.debug.tiny_random import initialize_sparse_model @@ -201,8 +190,9 @@ def __init__( show_progress=self.parallel_context.world_rank == 0, progress_rank=0 if self.parallel_context.world_rank == 0 else None, ) - if hf_config.model_type in {"qwen3_moe", "minimax_m2"}: - self.model.warmup_moe() + warmup_moe = getattr(self.model, "warmup_moe", None) + if callable(warmup_moe): + warmup_moe() self.sampler = Sampler() @@ -566,6 +556,10 @@ def reset_after_warmup(self) -> None: if os.getenv("SPARSEVLLM_DELTAKV_CLEAR_ATTN_SCORE_BUFFERS_AFTER_WARMUP", "0") == "1": self.sparse_controller.clear_decode_attn_score_buffers() + def log_operator_implementations(self) -> None: + if self.parallel_context.world_rank == 0: + operator_registry.log_operator_implementations() + def warmup_moe_workspace(self, num_tokens: int) -> None: warmup_moe = getattr(self.model, "warmup_moe", None) if not callable(warmup_moe): @@ -1038,6 +1032,8 @@ def prepare_sample(self, seqs: list[Sequence]): ) def _auto_capture_greedy_sampling(self, seqs: list[Sequence]) -> bool: + if self._uses_dpa_ep() and self.parallel_context.dp_size > 1: + return False if any(self._has_sampling_penalty(seq) for seq in seqs): return False if self.config.decode_cuda_graph_capture_sampling: @@ -1054,6 +1050,71 @@ def _auto_capture_greedy_sampling(self, seqs: list[Sequence]) -> bool: for seq in seqs ) + def _uses_dpa_ep(self) -> bool: + topology = getattr(self.config, "parallel_topology", None) + return bool( + getattr( + topology, + "is_dpa_ep", + getattr(self.config, "uses_dpa_ep_layout", False), + ) + ) + + def _deepseek_v4_decode_partition( + self, + seqs: list[Sequence], + ) -> tuple[list[Sequence], int]: + """Use stable DPA owners while keeping one graph shape on every rank.""" + dp_size = int(self.parallel_context.dp_size) + dp_rank = int(self.parallel_context.dp_rank) + if dp_size == 1: + return list(seqs), len(seqs) + owner_counts = [0] * dp_size + for seq in seqs: + owner_counts[int(seq.seq_id) % dp_size] += 1 + target = max(owner_counts) + owned = [seq for seq in seqs if int(seq.seq_id) % dp_size == dp_rank] + if len(owned) == target: + return owned, len(owned) + fillers = [seq for seq in seqs if int(seq.seq_id) % dp_size != dp_rank] + needed = target - len(owned) + if len(fillers) < needed: + raise RuntimeError( + "DeepSeek V4 DPA could not balance decode ranks: " + f"rank={dp_rank} owned={len(owned)} target={target} " + f"fillers={len(fillers)}." + ) + return owned + fillers[:needed], len(owned) + + def _gather_deepseek_v4_logits( + self, + local_logits: torch.Tensor, + seqs: list[Sequence], + local_owned: int, + ) -> torch.Tensor | None: + dp_size = int(self.parallel_context.dp_size) + if dp_size == 1: + return local_logits + target = int(local_logits.shape[0]) + if local_owned > target: + raise RuntimeError( + f"DeepSeek V4 local owned batch {local_owned} exceeds " + f"padded batch {target}." + ) + owned_logits = torch.zeros_like(local_logits) + owned_logits[:local_owned].copy_(local_logits[:local_owned]) + gathered = self.parallel_context.ep_all_gather_into_tensor(owned_logits) + if self.rank != 0: + return None + gathered = gathered.view(dp_size, target, -1) + owner_offsets = [0] * dp_size + ordered = [] + for seq in seqs: + owner = int(seq.seq_id) % dp_size + ordered.append(gathered[owner, owner_offsets[owner]]) + owner_offsets[owner] += 1 + return torch.stack(ordered, dim=0) + @staticmethod def _has_sampling_penalty(seq: Sequence) -> bool: return ( @@ -1249,28 +1310,48 @@ def run( with profiler.record(name): if not is_prefill: try: + global_seqs = seqs + run_seqs, local_owned = ( + self._deepseek_v4_decode_partition(seqs) + if self._uses_dpa_ep() + else (seqs, len(seqs)) + ) if self.config.decode_cuda_graph: logits, graph_token_ids = self.decode_cuda_graph_runner.run( - seqs, - capture_sampling=self._auto_capture_greedy_sampling(seqs), + run_seqs, + capture_sampling=self._auto_capture_greedy_sampling(run_seqs), ) else: - logits = self.decode_cuda_graph_runner.run_eager_static(seqs) + logits = self.decode_cuda_graph_runner.run_eager_static(run_seqs) + graph_token_ids = None + if self._uses_dpa_ep(): + logits = self._gather_deepseek_v4_logits( + logits, + global_seqs, + local_owned, + ) graph_token_ids = None if self.rank != 0: - self._post_sparse_forward(seqs, is_prefill) + self._post_sparse_forward(run_seqs, is_prefill) return None, None - self._post_sparse_forward(seqs, is_prefill) + self._post_sparse_forward(run_seqs, is_prefill) with profiler.record("model_sampler"): - sampling_logits = self._apply_sampling_penalties(logits, seqs) + sampling_logits = self._apply_sampling_penalties( + logits, + global_seqs, + ) token_ids = self._sample_model_outputs( sampling_logits, - seqs, + global_seqs, graph_token_ids=graph_token_ids, ) logprob_outputs = self._mask_recompute_logprobs( - seqs, - self._collect_logprobs(sampling_logits, token_ids, seqs), + global_seqs, + self._collect_logprobs( + sampling_logits, + token_ids, + global_seqs, + ), ) return token_ids, logprob_outputs finally: diff --git a/src/sparsevllm/engine/recurrent_state_manager.py b/src/sparsevllm/engine/recurrent_state_manager.py index 3623697a..a7609bb6 100644 --- a/src/sparsevllm/engine/recurrent_state_manager.py +++ b/src/sparsevllm/engine/recurrent_state_manager.py @@ -360,13 +360,36 @@ def get_decode_layer_state( f"layer_idx={layer_idx}; missing={missing}." ) state_buffers = {name: buffers[name] for name in self.state_spec.state_names} + tensor_specs = { + tensor_spec.name: tensor_spec + for tensor_spec in self.state_spec.tensor_specs + } for name, buffer in state_buffers.items(): - if buffer.dtype != dtype or buffer.device != device: + expected_dtype = tensor_specs[name].dtype + if buffer.dtype != expected_dtype or buffer.device != device: raise RuntimeError( - f"{self.state_spec.name} decode recurrent pool dtype/device does not match activations: " + f"{self.state_spec.name} decode recurrent pool does not match its model schema: " f"layer_idx={layer_idx} name={name} pool={buffer.dtype}/{buffer.device} " - f"expected={dtype}/{device}." + f"expected={expected_dtype}/{device}." ) + activation_state_name = ( + "conv_state" + if "conv_state" in tensor_specs + else ( + self.state_spec.state_names[0] + if len({spec.dtype for spec in self.state_spec.tensor_specs}) == 1 + else None + ) + ) + if ( + activation_state_name is not None + and tensor_specs[activation_state_name].dtype != dtype + ): + raise RuntimeError( + f"{self.state_spec.name} decode activation dtype does not match " + f"{activation_state_name}: activation={dtype} " + f"state={tensor_specs[activation_state_name].dtype}." + ) expected_rows: list[int] = [] for seq in seqs: row = self.seq_id_to_row.get(int(seq.seq_id)) diff --git a/src/sparsevllm/layers/attention_backend.py b/src/sparsevllm/layers/attention_backend.py index 72aacffd..6042cf74 100644 --- a/src/sparsevllm/layers/attention_backend.py +++ b/src/sparsevllm/layers/attention_backend.py @@ -3,6 +3,7 @@ import torch from sparsevllm.engine.cache_manager import DecodeComputeView, PrefillComputeView +from sparsevllm.operators.registry import record_operator_binding from sparsevllm.utils.context import get_context from sparsevllm.triton_kernel.context_flashattention_nopad import context_attention_fwd from sparsevllm.triton_kernel.flash_decoding_stage1 import flash_decode_stage1 as mha_flash_decode_stage1 @@ -94,6 +95,11 @@ def _fill_fake_attention_score(attn_score: torch.Tensor | None) -> None: class TritonAttentionBackend: """Thin backend wrapper around the existing Sparse-vLLM Triton attention kernels.""" + name = "triton" + + def __init__(self) -> None: + record_operator_binding("Attention", self) + def run_prefill( self, q: torch.Tensor, diff --git a/src/sparsevllm/layers/layernorm.py b/src/sparsevllm/layers/layernorm.py index 824347ab..c4efbc25 100755 --- a/src/sparsevllm/layers/layernorm.py +++ b/src/sparsevllm/layers/layernorm.py @@ -10,6 +10,8 @@ import torch from torch import nn +from sparsevllm.operators.registry import record_operator_binding + RMSNormFn = Callable[[torch.Tensor, torch.Tensor, float], torch.Tensor] FusedAddRMSNormFn = Callable[ @@ -140,6 +142,7 @@ def __init__( zero_centered_weight=self.zero_centered_weight, provider=provider, ) + record_operator_binding("RMSNorm", self) @property def provider_name(self) -> str: diff --git a/src/sparsevllm/layers/linear.py b/src/sparsevllm/layers/linear.py index efaee305..018c9a03 100755 --- a/src/sparsevllm/layers/linear.py +++ b/src/sparsevllm/layers/linear.py @@ -403,9 +403,11 @@ def __init__( output_size: int, bias: bool = False, quantization=None, + reduce_results: bool = True, ): tp_size = get_parallel_context().tp_size super().__init__(divide(input_size, tp_size), output_size, bias, 1, quantization=quantization) + self.reduce_results = bool(reduce_results) def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): param_data = param.data @@ -452,4 +454,4 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y = self.quant_provider(x, self.weight, self.weight_scale_inv, bias) else: y = F.linear(x, self.weight, bias) - return self.parallel_context.tp_all_reduce(y) + return self.parallel_context.tp_all_reduce(y) if self.reduce_results else y diff --git a/src/sparsevllm/method_registry.py b/src/sparsevllm/method_registry.py index 09aff5a2..e0758886 100644 --- a/src/sparsevllm/method_registry.py +++ b/src/sparsevllm/method_registry.py @@ -2,6 +2,9 @@ from dataclasses import dataclass +from sparsevllm.distributed.topology import ParallelMode, ParallelTopology +from sparsevllm.models.spec import MODEL_SPECS + PREFILL_POLICY_ALL_CHUNKED = "all_chunked" PREFILL_POLICY_LONG_BS1FULL_SHORT_BATCH = "long_bs1full_short_batch" PREFILL_POLICY_AUTO = "auto" @@ -56,123 +59,77 @@ "skipkv", } -H2O_SUPPORTED_MODEL_TYPES = frozenset( +SKIPKV_ASSET_MODEL_NAMES = frozenset( { - "qwen2", - "qwen3", - "qwen3_moe", - "qwen3_5", - "llama", - "minimax_m2", + "DeepSeek-R1-Distill-Llama-8B", + "DeepSeek-R1-Distill-Qwen-7B", + "DeepSeek-R1-Distill-Qwen-14B", } ) @dataclass(frozen=True) class ModelRuntimeCompatibility: - parallel_mode: str sparse_methods: frozenset[str] prefix_cache_methods: frozenset[str] requires_eager: bool = True decode_cuda_graph_methods: frozenset[str] = frozenset() +_MOE_SPARSE_METHODS = frozenset( + {"", "streamingllm", "snapkv", "h2o", "pyramidkv", "omnikv", "quest", "rkv"} +) + +DENSE_MODEL_COMPATIBILITY = ModelRuntimeCompatibility( + sparse_methods=frozenset(CANONICAL_SPARSE_METHODS), + prefix_cache_methods=frozenset(PREFIX_CACHE_SUPPORTED_METHODS), + requires_eager=False, + decode_cuda_graph_methods=frozenset(CANONICAL_SPARSE_METHODS), +) + + QWEN3_MOE_EP_COMPATIBILITY = ModelRuntimeCompatibility( - parallel_mode="ep_replicated_kv", - sparse_methods=frozenset( - { - "", - "streamingllm", - "snapkv", - "h2o", - "pyramidkv", - "omnikv", - "quest", - "rkv", - } - ), + sparse_methods=_MOE_SPARSE_METHODS, prefix_cache_methods=frozenset( {"", "omnikv", "quest", "snapkv", "h2o", "pyramidkv", "rkv"} ), requires_eager=False, - decode_cuda_graph_methods=frozenset( - { - "", - "streamingllm", - "snapkv", - "h2o", - "pyramidkv", - "omnikv", - "quest", - "rkv", - } - ), + decode_cuda_graph_methods=_MOE_SPARSE_METHODS, ) QWEN3_MOE_TP_EP_COMPATIBILITY = ModelRuntimeCompatibility( - parallel_mode="outer_tp_moe_tp_ep", - sparse_methods=frozenset( - { - "", - "streamingllm", - "snapkv", - "h2o", - "pyramidkv", - "omnikv", - "quest", - "rkv", - } - ), + sparse_methods=_MOE_SPARSE_METHODS, prefix_cache_methods=frozenset({""}), requires_eager=False, - decode_cuda_graph_methods=frozenset( - { - "", - "streamingllm", - "snapkv", - "h2o", - "pyramidkv", - "omnikv", - "quest", - "rkv", - } - ), + decode_cuda_graph_methods=_MOE_SPARSE_METHODS, ) QWEN3_MOE_TP_COMPATIBILITY = QWEN3_MOE_TP_EP_COMPATIBILITY -MINIMAX_M2_EP_COMPATIBILITY = ModelRuntimeCompatibility( - parallel_mode="ep_replicated_kv", - sparse_methods=frozenset( - { - "", - "streamingllm", - "snapkv", - "h2o", - "pyramidkv", - "omnikv", - "quest", - "rkv", - } +QWEN35_MOE_COMPATIBILITY = ModelRuntimeCompatibility( + sparse_methods=QWEN3_MOE_TP_EP_COMPATIBILITY.sparse_methods, + prefix_cache_methods=frozenset({""}), + requires_eager=False, + decode_cuda_graph_methods=( + QWEN3_MOE_TP_EP_COMPATIBILITY.decode_cuda_graph_methods ), +) + +DEEPSEEK_V4_DPA_EP_COMPATIBILITY = ModelRuntimeCompatibility( + sparse_methods=frozenset({""}), + prefix_cache_methods=frozenset(), + requires_eager=False, + decode_cuda_graph_methods=frozenset({""}), +) + +MINIMAX_M2_EP_COMPATIBILITY = ModelRuntimeCompatibility( + sparse_methods=_MOE_SPARSE_METHODS, prefix_cache_methods=frozenset({"", "omnikv", "quest"}), requires_eager=False, - decode_cuda_graph_methods=frozenset( - { - "", - "streamingllm", - "snapkv", - "h2o", - "pyramidkv", - "omnikv", - "quest", - "rkv", - } - ), + decode_cuda_graph_methods=_MOE_SPARSE_METHODS, ) MINIMAX_M2_TP_EP_COMPATIBILITY = ModelRuntimeCompatibility( - parallel_mode="outer_tp_moe_tp_ep", sparse_methods=MINIMAX_M2_EP_COMPATIBILITY.sparse_methods, prefix_cache_methods=MINIMAX_M2_EP_COMPATIBILITY.prefix_cache_methods, requires_eager=False, @@ -180,10 +137,25 @@ class ModelRuntimeCompatibility: ) MODEL_RUNTIME_COMPATIBILITY = { - "qwen3_moe": QWEN3_MOE_EP_COMPATIBILITY, - "minimax_m2": MINIMAX_M2_EP_COMPATIBILITY, + **{ + (model_type, ParallelMode.STANDARD): DENSE_MODEL_COMPATIBILITY + for model_type in ("qwen2", "qwen3", "qwen3_5", "llama") + }, + ("qwen3_moe", ParallelMode.STANDARD): QWEN3_MOE_EP_COMPATIBILITY, + ("qwen3_moe", ParallelMode.OUTER_TP_MOE): QWEN3_MOE_TP_EP_COMPATIBILITY, + ("qwen3_5_moe", ParallelMode.STANDARD): QWEN35_MOE_COMPATIBILITY, + ("qwen3_5_moe", ParallelMode.OUTER_TP_MOE): QWEN35_MOE_COMPATIBILITY, + ("minimax_m2", ParallelMode.STANDARD): MINIMAX_M2_EP_COMPATIBILITY, + ("minimax_m2", ParallelMode.OUTER_TP_MOE): MINIMAX_M2_TP_EP_COMPATIBILITY, + ("deepseek_v4", ParallelMode.DPA_EP): DEEPSEEK_V4_DPA_EP_COMPATIBILITY, } +H2O_SUPPORTED_MODEL_TYPES = frozenset( + model_type + for (model_type, _parallel_mode), compatibility in MODEL_RUNTIME_COMPATIBILITY.items() + if "h2o" in compatibility.sparse_methods +) + # All shipped cache managers now expose a graph-stable decode preparation path. DECODE_CUDA_GRAPH_SUPPORTED_METHODS = set(CANONICAL_SPARSE_METHODS) TP_DECODE_CUDA_GRAPH_SUPPORTED_METHODS = { @@ -234,6 +206,18 @@ def normalize_sparse_method(method: str | None) -> str: return METHOD_ALIASES.get(normalized, normalized) +def validate_sparse_method_assets(method: str | None, model_path: str) -> None: + if normalize_sparse_method(method) != "skipkv": + return + model_name = str(model_path).rstrip("/").split("/")[-1] + if model_name not in SKIPKV_ASSET_MODEL_NAMES: + supported = ", ".join(sorted(SKIPKV_ASSET_MODEL_NAMES)) + raise ValueError( + "SkipKV is supported only for models with released steering assets: " + f"{supported}. Got model basename {model_name!r}." + ) + + def is_deltakv_method(method: str | None) -> bool: return normalize_sparse_method(method) == "deltakv" @@ -250,42 +234,20 @@ def validate_model_runtime_compatibility( *, model_type: str, sparse_method: str | None, - tensor_parallel_size: int, - expert_parallel_size: int, - data_parallel_size: int, + topology: ParallelTopology, enforce_eager: bool, decode_cuda_graph: bool, enable_prefix_caching: bool, -) -> ModelRuntimeCompatibility | None: +) -> ModelRuntimeCompatibility: model_type = str(model_type or "").strip().lower() - compatibility = MODEL_RUNTIME_COMPATIBILITY.get(model_type) - if compatibility is None: - return None - method = normalize_sparse_method(sparse_method) - tp_size = int(tensor_parallel_size) - ep_size = int(expert_parallel_size) - dp_size = int(data_parallel_size) - if model_type in {"qwen3_moe", "minimax_m2"} and tp_size > 1: - compatibility = ( - QWEN3_MOE_TP_EP_COMPATIBILITY - if model_type == "qwen3_moe" - else MINIMAX_M2_TP_EP_COMPATIBILITY - ) - if dp_size != 1: - raise ValueError( - f"{model_type} outer_tp_moe_tp_ep requires DP=1, got " - f"TP={tp_size}, EP={ep_size}, DP={dp_size}." - ) - elif tp_size != 1 or dp_size != 1: - raise ValueError( - f"{model_type} {compatibility.parallel_mode} requires TP=1 and DP=1, got " - f"TP={tensor_parallel_size}, EP={expert_parallel_size}, DP={data_parallel_size}." - ) - if ep_size <= 0: - raise ValueError( - f"{model_type} requires a positive expert_parallel_size, got {expert_parallel_size}." + compatibility = MODEL_RUNTIME_COMPATIBILITY.get((model_type, topology.mode)) + if compatibility is None: + raise NotImplementedError( + f"Unsupported Sparse-vLLM model_type={model_type!r} with " + f"parallel mode={topology.mode.value!r}." ) + if compatibility.requires_eager and not bool(enforce_eager): raise ValueError(f"{model_type} v1 requires enforce_eager=True.") if bool(decode_cuda_graph) and method not in compatibility.decode_cuda_graph_methods: @@ -297,22 +259,13 @@ def validate_model_runtime_compatibility( f"{model_type} v1 decode_cuda_graph is validated only for {supported}; " f"got method={method!r}." ) - if model_type == "qwen3_moe" and method == "skipkv": - raise NotImplementedError( - "Qwen3MoE + SkipKV requires a Qwen3MoE-matched steering asset and validation; " - "no compatible asset is currently registered." - ) - if model_type == "qwen3_moe" and method == "deltakv": - raise NotImplementedError( - "Qwen3MoE + DeltaKV is not part of the validated v1 compatibility matrix." - ) if method not in compatibility.sparse_methods: supported = ", ".join( "'vanilla'" if item == "" else repr(item) for item in sorted(compatibility.sparse_methods) ) raise ValueError( - f"Unsupported {model_type} {compatibility.parallel_mode} sparse method " + f"Unsupported {model_type} {topology.mode.value} sparse method " f"{method!r}; validated methods: {supported}." ) if bool(enable_prefix_caching) and method not in compatibility.prefix_cache_methods: diff --git a/src/sparsevllm/models/checkpoint.py b/src/sparsevllm/models/checkpoint.py new file mode 100644 index 00000000..a8d545e4 --- /dev/null +++ b/src/sparsevllm/models/checkpoint.py @@ -0,0 +1,348 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch + +from sparsevllm.distributed.topology import ParallelTopology +from sparsevllm.quantization.config import QuantizationConfig +from sparsevllm.utils.config import config_get + + +def _validate_architecture(model_name: str, config: Any, expected: str) -> None: + architectures = tuple(config_get(config, "architectures", ()) or ()) + if architectures != (expected,): + raise ValueError( + f"{model_name} requires architectures=[{expected!r}], " + f"got {list(architectures)}." + ) + + +def _validate_bf16(model_name: str, config: Any, description: str) -> None: + dtype = config_get(config, "torch_dtype", None) + if dtype is None: + dtype = config_get(config, "dtype", None) + if dtype not in {torch.bfloat16, "bfloat16"}: + raise ValueError(f"{model_name} requires {description}, got dtype={dtype!r}.") + + +def _validate_fields( + model_name: str, + config: Any, + expected_fields: Mapping[str, Any], +) -> None: + for field, expected in expected_fields.items(): + actual = config_get(config, field, None) + if actual != expected: + raise ValueError( + f"{model_name} requires {field}={expected!r}, got {actual!r}." + ) + + +def _validate_exclusions( + model_name: str, + raw_quantization_config: Any, + required: set[str], + description: str, +) -> None: + excluded = { + str(name) + for name in ( + config_get(raw_quantization_config, "modules_to_not_convert", ()) or () + ) + } + missing = sorted(required - excluded) + if missing: + raise ValueError( + f"{model_name} quantization_config must exclude {description}; " + f"missing {missing[:8]}." + ) + + +def _validate_qwen3_fp8(config: Any, topology: ParallelTopology) -> None: + _validate_architecture("Qwen3 FP8", config, "Qwen3ForCausalLM") + _validate_bf16("Qwen3 FP8", config, "BF16 non-quantized parameters") + head_dim = int(config_get(config, "head_dim", 0) or 0) + dimensions = { + "hidden_size": int(config_get(config, "hidden_size", 0) or 0), + "intermediate_size": int(config_get(config, "intermediate_size", 0) or 0), + "query_size": int(config_get(config, "num_attention_heads", 0) or 0) + * head_dim, + "key_value_size": int( + config_get(config, "num_key_value_heads", 0) or 0 + ) + * head_dim, + } + alignment = 128 * topology.attention_tp_size + invalid = { + name: size + for name, size in dimensions.items() + if size <= 0 or size % alignment + } + if invalid: + raise ValueError( + "Qwen3 FP8 requires every TP-local dense projection dimension to be " + f"128-aligned; TP={topology.attention_tp_size}, invalid={invalid}." + ) + + +def _validate_qwen3_moe_fp8(config: Any, raw_quantization_config: Any) -> None: + _validate_architecture("Qwen3MoE FP8", config, "Qwen3MoeForCausalLM") + _validate_bf16("Qwen3MoE FP8", config, "BF16 non-quantized parameters") + dimensions = { + "hidden_size": int(config_get(config, "hidden_size", 0) or 0), + "moe_intermediate_size": int( + config_get(config, "moe_intermediate_size", 0) or 0 + ), + } + invalid = { + name: size for name, size in dimensions.items() if size <= 0 or size % 128 + } + if invalid: + raise ValueError( + "Qwen3MoE FP8 requires hidden_size and moe_intermediate_size aligned " + f"to 128, invalid={invalid}." + ) + required = {"lm_head"} + required.update( + f"model.layers.{layer_idx}.mlp.gate" + for layer_idx in range(int(config_get(config, "num_hidden_layers", 0) or 0)) + ) + _validate_exclusions( + "Qwen3MoE FP8", + raw_quantization_config, + required, + "lm_head and every router gate", + ) + + +def _validate_qwen3_moe( + config: Any, + raw_quantization_config: Any, + quantization: QuantizationConfig, + topology: ParallelTopology, +) -> None: + decoder_sparse_step = int(config_get(config, "decoder_sparse_step", 1)) + mlp_only_layers = tuple( + int(layer_idx) + for layer_idx in (config_get(config, "mlp_only_layers", ()) or ()) + ) + if decoder_sparse_step != 1 or mlp_only_layers: + raise NotImplementedError( + "Qwen3MoE v1 requires every decoder layer to be MoE, got " + f"decoder_sparse_step={decoder_sparse_step}, " + f"mlp_only_layers={list(mlp_only_layers)}." + ) + shared_intermediate_size = int( + config_get(config, "shared_expert_intermediate_size", 0) or 0 + ) + if shared_intermediate_size: + raise NotImplementedError( + "Qwen3MoE v1 does not support shared experts, got " + f"shared_expert_intermediate_size={shared_intermediate_size}." + ) + dtype = config_get(config, "torch_dtype", None) + if dtype not in {torch.bfloat16, torch.float16}: + raise NotImplementedError( + "Qwen3MoE v1 supports BF16/FP16 checkpoints only, " + f"got torch_dtype={dtype}." + ) + if topology.is_outer_tp_moe and dtype != torch.bfloat16: + raise NotImplementedError( + "Qwen3MoE outer TP supports BF16 checkpoints only, " + f"got torch_dtype={dtype}." + ) + if quantization.enabled: + _validate_qwen3_moe_fp8(config, raw_quantization_config) + + +def _validate_qwen35_moe( + outer_config: Any, + config: Any, + quantization: QuantizationConfig, + topology: ParallelTopology, +) -> None: + _validate_architecture( + "Qwen3.6 MoE", outer_config, "Qwen3_5MoeForConditionalGeneration" + ) + _validate_bf16( + "Qwen3.6 MoE", + config, + "BF16 activations with either BF16 or block-FP8 language-model weights", + ) + _validate_fields( + "Qwen3.6 MoE", + config, + { + "hidden_act": "silu", + "attn_output_gate": True, + "attention_bias": False, + "partial_rotary_factor": 0.25, + "mamba_ssm_dtype": "float32", + "rms_norm_eps": 1.0e-6, + "tie_word_embeddings": False, + }, + ) + if quantization.enabled: + dimensions = { + "hidden_size": int(config_get(config, "hidden_size", 0) or 0), + "shared_expert_intermediate_size": int( + config_get(config, "shared_expert_intermediate_size", 0) or 0 + ) + // topology.attention_tp_size, + } + invalid = { + name: size + for name, size in dimensions.items() + if size <= 0 or size % 128 + } + if invalid: + raise ValueError( + "Qwen3.6 MoE FP8 local Linear dimensions must be 128-aligned, " + f"got TP={topology.attention_tp_size}, invalid={invalid}." + ) + + +def _validate_minimax(config: Any, raw_quantization_config: Any) -> None: + _validate_architecture("MiniMax M2.7", config, "MiniMaxM2ForCausalLM") + _validate_fields( + "MiniMax M2.7", + config, + { + "hidden_act": "silu", + "qk_norm_type": "per_layer", + "scoring_func": "sigmoid", + "use_qk_norm": True, + "use_routing_bias": True, + "use_mtp": True, + "tie_word_embeddings": False, + }, + ) + _validate_bf16("MiniMax M2.7", config, "BF16 non-quantized parameters") + _validate_exclusions( + "MiniMax M2.7", + raw_quantization_config, + {"gate", "e_score_correction_bias", "lm_head"}, + "gate, e_score_correction_bias, and lm_head", + ) + + +_DEEPSEEK_V4_FLASH_FIXED_FIELDS = { + "hidden_size": 4096, + "head_dim": 512, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "q_lora_rank": 1024, + "o_lora_rank": 1024, + "o_groups": 8, + "qk_rope_head_dim": 64, + "index_n_heads": 64, + "index_head_dim": 128, + "index_topk": 512, + "n_routed_experts": 256, + "num_experts_per_tok": 6, + "hc_mult": 4, + "sliding_window": 128, +} + + +def _validate_deepseek_v4(config: Any, raw_quantization_config: Any) -> None: + _validate_architecture( + "DeepSeek V4 Flash", config, "DeepseekV4ForCausalLM" + ) + _validate_bf16( + "DeepSeek V4 Flash", config, "BF16 non-quantized parameters" + ) + _validate_fields( + "DeepSeek V4 Flash", + config, + _DEEPSEEK_V4_FLASH_FIXED_FIELDS, + ) + expected_compress_rates = { + "compressed_sparse_attention": 4, + "heavily_compressed_attention": 128, + } + if config_get(config, "compress_rates", {}) != expected_compress_rates: + raise ValueError( + "DeepSeek V4 Flash requires CSA/HCA compression rates 4/128, got " + f"{config_get(config, 'compress_rates', {})!r}." + ) + _validate_fields( + "DeepSeek V4 Flash", + config, + { + "num_nextn_predict_layers": 1, + "expert_dtype": "fp4", + }, + ) + expected_quantization = { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + } + for field, expected in expected_quantization.items(): + actual = config_get(raw_quantization_config, field, None) + if actual != expected: + raise ValueError( + "DeepSeek V4 Flash quantization_config requires " + f"{field}={expected!r}, got {actual!r}." + ) + + +def validate_checkpoint( + model_type: str, + *, + outer_config: Any, + config: Any, + raw_quantization_config: Any, + quantization: QuantizationConfig, + topology: ParallelTopology, +) -> None: + validator = CHECKPOINT_VALIDATORS.get(model_type) + if validator is not None: + validator( + outer_config, + config, + raw_quantization_config, + quantization, + topology, + ) + + +def _qwen35_checkpoint(_outer, config, _raw, quantization, _topology) -> None: + if not quantization.enabled: + _validate_bf16("Qwen3.5", config, "BF16 weights") + + +def _qwen35_moe_checkpoint(outer, config, _raw, quantization, topology) -> None: + _validate_qwen35_moe(outer, config, quantization, topology) + + +def _minimax_checkpoint(_outer, config, raw, _quantization, _topology) -> None: + _validate_minimax(config, raw) + + +def _qwen3_checkpoint(_outer, config, _raw, quantization, topology) -> None: + if quantization.enabled: + _validate_qwen3_fp8(config, topology) + + +def _qwen3_moe_checkpoint(_outer, config, raw, quantization, topology) -> None: + _validate_qwen3_moe(config, raw, quantization, topology) + + +def _deepseek_v4_checkpoint(_outer, config, raw, _quantization, _topology) -> None: + _validate_deepseek_v4(config, raw) + + +CHECKPOINT_VALIDATORS = { + "qwen3": _qwen3_checkpoint, + "qwen3_moe": _qwen3_moe_checkpoint, + "qwen3_5": _qwen35_checkpoint, + "qwen3_5_moe": _qwen35_moe_checkpoint, + "minimax_m2": _minimax_checkpoint, + "deepseek_v4": _deepseek_v4_checkpoint, +} diff --git a/src/sparsevllm/models/deepseek_v4.py b/src/sparsevllm/models/deepseek_v4.py new file mode 100644 index 00000000..acbdbd29 --- /dev/null +++ b/src/sparsevllm/models/deepseek_v4.py @@ -0,0 +1,293 @@ +from __future__ import annotations + +import re + +import torch +from torch import nn +from transformers import DynamicCache +from transformers.models.deepseek_v4.modeling_deepseek_v4 import DeepseekV4Model + +from sparsevllm.layers.embed_head import ParallelLMHead +from sparsevllm.models.deepseek_v4_native import DeepseekV4Model as NativeDeepseekV4Model +from sparsevllm.utils.context import get_context +from sparsevllm.utils.weight_target import WeightTarget + + +_EXPERT_SOURCE_RE = re.compile( + r"^layers\.(\d+)\.ffn\.experts\.(\d+)\.(w1|w2|w3)\.weight$" +) +_EXPERT_TARGET_RE = re.compile( + r"^model\.layers\.(\d+)\.ffn\.experts\.(\d+)\.(gate|down|up)\.expert_weight$" +) + + +class DeepseekV4ForCausalLM(nn.Module): + """DeepSeek V4 architecture adapter used by the tiny-random reference path. + + The formal FP4 DPA+EP runtime replaces the Transformers layers with native + Sparse-vLLM operators. Keeping this reference adapter exact gives that path + a deterministic end-to-end correctness oracle without reading checkpoint + tensors during development. + """ + + special_weight_loaders = (".expert_weight",) + + def __init__(self, config) -> None: + super().__init__() + self.config = config + self.tiny_random = bool(getattr(config, "sparsevllm_tiny_random", False)) + self.model = DeepseekV4Model(config) if self.tiny_random else NativeDeepseekV4Model(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + self._seq_caches: dict[int, DynamicCache] = {} + self._skipped_expert_weights: set[str] = set() + self._skipped_expert_scales: set[str] = set() + + @staticmethod + def scale_key_for_weight(weight_key: str) -> str | None: + if not weight_key.endswith(".weight"): + return None + return weight_key[: -len(".weight")] + ".scale" + + @staticmethod + def weight_key_for_scale(scale_key: str) -> str | None: + if not scale_key.endswith(".scale"): + return None + return scale_key[: -len(".scale")] + ".weight" + + def map_weight_name(self, source_weight_name: str) -> str | None: + expert_match = _EXPERT_SOURCE_RE.match(source_weight_name) + if expert_match is not None: + layer_idx, expert_id, source_projection = expert_match.groups() + experts = self.model.layers[int(layer_idx)].ffn.experts + if not experts.is_local_expert(int(expert_id)): + self._skipped_expert_weights.add(source_weight_name) + return None + projection = {"w1": "gate", "w2": "down", "w3": "up"}[ + source_projection + ] + return ( + f"model.layers.{layer_idx}.ffn.experts.{expert_id}." + f"{projection}.expert_weight" + ) + if source_weight_name.startswith(("mtp.", "model.mtp.", "nextn.")): + return None + if source_weight_name.startswith("embed."): + mapped = "model." + source_weight_name + elif source_weight_name.startswith("layers."): + mapped = "model." + source_weight_name + elif source_weight_name.startswith("norm."): + mapped = "model." + source_weight_name + elif source_weight_name.startswith("head."): + mapped = "lm_" + source_weight_name + elif source_weight_name.startswith("hc_head_"): + mapped = "model." + source_weight_name + else: + mapped = source_weight_name + replacements = ( + (".attn.wq_a.", ".attn.q_a_proj."), + (".attn.wq_b.", ".attn.q_b_proj."), + (".attn.wkv.", ".attn.kv_proj."), + (".attn.wo_a.", ".attn.o_a_proj."), + (".attn.wo_b.", ".attn.o_b_proj."), + (".attn.q_norm.", ".attn.q_a_norm."), + (".attn.attn_sink", ".attn.sinks"), + (".attn.compressor.wkv.", ".attn.compressor.kv_proj."), + (".attn.compressor.wgate.", ".attn.compressor.gate_proj."), + (".attn.compressor.ape", ".attn.compressor.position_bias"), + (".attn.compressor.norm.", ".attn.compressor.kv_norm."), + ( + ".attn.indexer.compressor.wkv.", + ".attn.compressor.indexer.kv_proj.", + ), + ( + ".attn.indexer.compressor.wgate.", + ".attn.compressor.indexer.gate_proj.", + ), + ( + ".attn.indexer.compressor.ape", + ".attn.compressor.indexer.position_bias", + ), + ( + ".attn.indexer.compressor.norm.", + ".attn.compressor.indexer.kv_norm.", + ), + ( + ".attn.indexer.weights_proj.", + ".attn.compressor.indexer.scorer.weights_proj.", + ), + ( + ".attn.indexer.wq_b.", + ".attn.compressor.indexer.q_b_proj.", + ), + (".hc_attn_fn", ".hc_attn.fn"), + (".hc_attn_base", ".hc_attn.base"), + (".hc_attn_scale", ".hc_attn.scale"), + (".hc_ffn_fn", ".hc_ffn.fn"), + (".hc_ffn_base", ".hc_ffn.base"), + (".hc_ffn_scale", ".hc_ffn.scale"), + (".hc_head_fn", ".hc_head.fn"), + (".hc_head_base", ".hc_head.base"), + (".hc_head_scale", ".hc_head.scale"), + ) + for source, target in replacements: + mapped = mapped.replace(source, target) + return mapped + + def resolve_special_weight(self, target_weight_name: str) -> WeightTarget | None: + match = _EXPERT_TARGET_RE.match(target_weight_name) + if match is None: + return None + layer_idx, expert_id, projection = match.groups() + return WeightTarget( + self.model.layers[int(layer_idx)].ffn.experts, + (int(expert_id), projection), + ) + + def load_special_weight( + self, + target_weight_name: str, + loaded_weight: torch.Tensor, + loaded_scale: torch.Tensor | None, + ) -> int: + target = self.resolve_special_weight(target_weight_name) + if target is None: + return 0 + expert_id, projection = target.shard_id + target.module.load_expert_weight( + expert_id, projection, loaded_weight, loaded_scale + ) + return 1 + + def record_skipped_weight( + self, + source_weight_name: str, + loaded_weight_shape, + loaded_weight_dtype, + loaded_scale_shape, + loaded_scale_dtype, + ) -> None: + del loaded_weight_shape, loaded_weight_dtype, loaded_scale_shape, loaded_scale_dtype + if _EXPERT_SOURCE_RE.match(source_weight_name) is None: + if source_weight_name.startswith(("mtp.", "model.mtp.", "nextn.")): + return + raise ValueError(f"Unexpectedly skipped DeepSeek V4 tensor {source_weight_name!r}.") + self._skipped_expert_weights.add(source_weight_name) + self._skipped_expert_scales.add( + source_weight_name[: -len(".weight")] + ".scale" + ) + + def _forward_sequence( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + seq_id: int, + ) -> torch.Tensor: + cache = self._seq_caches.get(seq_id) + if cache is None and self.tiny_random: + if int(positions[0]) != 0: + raise RuntimeError( + f"DeepSeek V4 sequence {seq_id} has no cache at position " + f"{int(positions[0])}." + ) + cache = DynamicCache(config=self.config) + self._seq_caches[seq_id] = cache + if self.tiny_random: + output = self.model( + input_ids=input_ids.unsqueeze(0), + position_ids=positions.unsqueeze(0), + past_key_values=cache, + use_cache=True, + return_dict=True, + ) + return output.last_hidden_state.squeeze(0) + raise RuntimeError("Native DeepSeek V4 uses cache-manager rows, not per-sequence caches.") + + def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + context = get_context() + seqs = list(context.seqs or ()) + if not seqs: + raise RuntimeError("DeepSeek V4 forward requires active engine sequences.") + if context.is_prefill: + cu_seqlens = context.cu_seqlens_q + if cu_seqlens is None or int(cu_seqlens.numel()) != len(seqs) + 1: + raise RuntimeError("DeepSeek V4 prefill requires one cu_seqlens entry per sequence.") + if not self.tiny_random: + rows = context.cache_manager.get_layer_batch_states(0).req_indices + outputs = [] + for index in range(len(seqs)): + start = int(cu_seqlens[index]) + end = int(cu_seqlens[index + 1]) + outputs.append( + self.model( + input_ids[start:end].unsqueeze(0), + positions[start:end].unsqueeze(0), + rows[index : index + 1], + ).squeeze(0) + ) + return torch.cat(outputs, dim=0) + outputs = [] + for index, seq in enumerate(seqs): + start = int(cu_seqlens[index]) + end = int(cu_seqlens[index + 1]) + outputs.append( + self._forward_sequence( + input_ids[start:end], + positions[start:end], + int(seq.seq_id), + ) + ) + return torch.cat(outputs, dim=0) + if self.tiny_random and int(input_ids.numel()) != len(seqs): + raise RuntimeError( + "DeepSeek V4 decode requires exactly one token per active sequence, " + f"got tokens={input_ids.numel()} sequences={len(seqs)}." + ) + if not self.tiny_random: + rows = context.cache_manager.get_layer_batch_states(0).req_indices + return self.model(input_ids.unsqueeze(1), positions.unsqueeze(1), rows).squeeze(1) + return torch.cat( + [ + self._forward_sequence( + input_ids[index : index + 1], + positions[index : index + 1], + int(seq.seq_id), + ) + for index, seq in enumerate(seqs) + ], + dim=0, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.lm_head(hidden_states) + + @torch.inference_mode() + def warmup_moe(self, num_tokens: int = 1) -> None: + if self.tiny_random: + return + for layer in self.model.layers: + layer.ffn.experts.prepare_for_inference() + + def validate_loaded_weights(self, loaded_parameter_names: set[str]) -> None: + if self.tiny_random: + return + expert_parameters = { + name + for name, _ in self.named_parameters() + if name.endswith(".ffn.experts.w13_weight") + or name.endswith(".ffn.experts.w2_weight") + } + missing = sorted( + {name for name, _ in self.named_parameters()} + - expert_parameters + - loaded_parameter_names + ) + if missing: + raise ValueError(f"Missing replicated DeepSeek V4 parameters: {missing[:8]}.") + for layer in self.model.layers: + layer.ffn.experts.validate_loaded_weights() + + def free_sequence_cache(self, seq_id: int) -> None: + self._seq_caches.pop(int(seq_id), None) + + def reset_after_warmup(self) -> None: + self._seq_caches.clear() diff --git a/src/sparsevllm/models/deepseek_v4_native.py b/src/sparsevllm/models/deepseek_v4_native.py new file mode 100644 index 00000000..28056f47 --- /dev/null +++ b/src/sparsevllm/models/deepseek_v4_native.py @@ -0,0 +1,676 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn + +from sparsevllm.distributed import get_parallel_context +from sparsevllm.layers.linear import ReplicatedLinear +from sparsevllm.operators.moe import MoeOpSpec, model_activation_dtype, resolve_moe_provider + +from transformers.models.deepseek_v4.modeling_deepseek_v4 import ( + DeepseekV4Attention as ReferenceDeepseekV4Attention, + DeepseekV4RMSNorm, + DeepseekV4RotaryEmbedding, + apply_rotary_pos_emb, +) +from sparsevllm.utils.context import get_context + + +class DeepseekV4GroupedFp8Linear(ReplicatedLinear): + """Eight independent grouped projections stored in one checkpoint tensor.""" + + def __init__(self, config) -> None: + self.num_groups = int(config.o_groups) + self.group_input_size = int(config.num_attention_heads * config.head_dim) // self.num_groups + self.group_output_size = int(config.o_lora_rank) + super().__init__( + self.group_input_size, + self.num_groups * self.group_output_size, + quantization=config.quantization_config, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.shape[-2:] != (self.num_groups, self.group_input_size): + raise ValueError( + "DeepSeek V4 grouped projection expects " + f"[..., {self.num_groups}, {self.group_input_size}], got {tuple(x.shape)}." + ) + if not self.quantized: + weight = self.weight.view( + self.num_groups, self.group_output_size, self.group_input_size + ) + return torch.einsum("...gi,goi->...go", x, weight) + outputs = [] + scale_rows = self.group_output_size // 128 + for group in range(self.num_groups): + row_start = group * self.group_output_size + scale_start = group * scale_rows + outputs.append( + self.quant_provider( + x[..., group, :], + self.weight[row_start : row_start + self.group_output_size], + self.weight_scale_inv[scale_start : scale_start + scale_rows], + ) + ) + return torch.stack(outputs, dim=-2) + + +class DeepseekV4HyperConnection(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.hc_mult = int(config.hc_mult) + self.sinkhorn_iters = int(config.hc_sinkhorn_iters) + self.eps = float(config.hc_eps) + mix = (2 + self.hc_mult) * self.hc_mult + self.fn = nn.Parameter(torch.empty(mix, self.hc_mult * int(config.hidden_size))) + self.base = nn.Parameter(torch.empty(mix)) + self.scale = nn.Parameter(torch.empty(3)) + self.rms_eps = float(config.rms_norm_eps) + + def forward(self, streams: torch.Tensor): + hc = self.hc_mult + flat = streams.flatten(start_dim=-2).float() + flat = flat * torch.rsqrt(flat.square().mean(-1, keepdim=True) + self.rms_eps) + pre_w, post_w, comb_w = F.linear(flat, self.fn.float()).split( + [hc, hc, hc * hc], dim=-1 + ) + pre_b, post_b, comb_b = self.base.float().split([hc, hc, hc * hc]) + pre_scale, post_scale, comb_scale = self.scale.float().unbind() + pre = torch.sigmoid(pre_w * pre_scale + pre_b) + self.eps + post = 2 * torch.sigmoid(post_w * post_scale + post_b) + comb = torch.softmax( + comb_w.view(*comb_w.shape[:-1], hc, hc) * comb_scale + + comb_b.view(hc, hc), + dim=-1, + ) + self.eps + comb = comb / (comb.sum(dim=-2, keepdim=True) + self.eps) + for _ in range(self.sinkhorn_iters - 1): + comb = comb / (comb.sum(dim=-1, keepdim=True) + self.eps) + comb = comb / (comb.sum(dim=-2, keepdim=True) + self.eps) + collapsed = (pre.unsqueeze(-1) * streams).sum(dim=-2).to(streams.dtype) + return post, comb, collapsed + + +class DeepseekV4HyperHead(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.hc_mult = int(config.hc_mult) + self.eps = float(config.hc_eps) + self.rms_eps = float(config.rms_norm_eps) + self.fn = nn.Parameter( + torch.empty(self.hc_mult, self.hc_mult * int(config.hidden_size)) + ) + self.base = nn.Parameter(torch.empty(self.hc_mult)) + self.scale = nn.Parameter(torch.empty(1)) + + def forward(self, streams: torch.Tensor) -> torch.Tensor: + flat = streams.flatten(start_dim=-2).float() + flat = flat * torch.rsqrt(flat.square().mean(-1, keepdim=True) + self.rms_eps) + weights = torch.sigmoid( + F.linear(flat, self.fn.float()) * self.scale.float() + self.base.float() + ) + self.eps + return (weights.unsqueeze(-1) * streams).sum(dim=-2).to(streams.dtype) + + +class DeepseekV4Router(nn.Module): + def __init__(self, config, *, hash_routing: bool) -> None: + super().__init__() + self.num_experts = int(config.n_routed_experts) + self.top_k = int(config.num_experts_per_tok) + self.hidden_size = int(config.hidden_size) + self.routed_scaling_factor = float(config.routed_scaling_factor) + self.hash_routing = bool(hash_routing) + self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_size)) + if self.hash_routing: + self.register_buffer( + "tid2eid", + torch.empty(int(config.vocab_size), self.top_k, dtype=torch.long), + ) + self.register_buffer("bias", None) + else: + self.register_buffer("tid2eid", None) + self.register_buffer("bias", torch.empty(self.num_experts)) + + def forward(self, hidden_states: torch.Tensor, input_ids: torch.Tensor): + scores = F.softplus(F.linear(hidden_states, self.weight)).sqrt() + if self.hash_routing: + topk_ids = self.tid2eid[input_ids].long() + else: + topk_ids = torch.topk( + scores + self.bias, + self.top_k, + dim=-1, + sorted=False, + ).indices + topk_weights = scores.gather(-1, topk_ids) + topk_weights = topk_weights / (topk_weights.sum(-1, keepdim=True) + 1e-20) + return topk_weights * self.routed_scaling_factor, topk_ids + + +class DeepseekV4SharedExperts(nn.Module): + def __init__(self, config) -> None: + super().__init__() + hidden_size = int(config.hidden_size) + intermediate_size = int(config.moe_intermediate_size) * int(config.n_shared_experts) + quantization = config.quantization_config + self.limit = float(config.swiglu_limit) + self.w1 = ReplicatedLinear(hidden_size, intermediate_size, quantization=quantization) + self.w2 = ReplicatedLinear(intermediate_size, hidden_size, quantization=quantization) + self.w3 = ReplicatedLinear(hidden_size, intermediate_size, quantization=quantization) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + gate = self.w1(hidden_states).clamp(max=self.limit) + up = self.w3(hidden_states).clamp(min=-self.limit, max=self.limit) + return self.w2(F.silu(gate) * up) + + +class DeepseekV4PackedExperts(nn.Module): + def __init__(self, config) -> None: + super().__init__() + parallel = get_parallel_context() + self.ep_rank = int(parallel.ep_rank) + self.ep_size = int(parallel.ep_size) + self.num_experts = int(config.n_routed_experts) + self.num_local_experts = self.num_experts // self.ep_size + self.local_expert_start = self.ep_rank * self.num_local_experts + self.local_expert_end = self.local_expert_start + self.num_local_experts + self.hidden_size = int(config.hidden_size) + self.intermediate_size = int(config.moe_intermediate_size) + self.op_spec = MoeOpSpec( + num_experts=self.num_experts, + num_local_experts=self.num_local_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + top_k=int(config.num_experts_per_tok), + activation_dtype=model_activation_dtype(config), + weight_dtype=torch.uint8, + block_shape=(1, 32), + ep_size=self.ep_size, + cuda_graph=bool(getattr(config, "decode_cuda_graph", False)), + scale_dtype=torch.uint8, + activation_limit=float(config.swiglu_limit), + ) + self.provider = resolve_moe_provider(self.op_spec) + self.w13_weight = nn.Parameter( + torch.empty( + self.num_local_experts, + 2 * self.intermediate_size, + self.hidden_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + self.w2_weight = nn.Parameter( + torch.empty( + self.num_local_experts, + self.hidden_size, + self.intermediate_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + self.register_buffer( + "w13_scale_inv", + torch.empty( + self.num_local_experts, + 2 * self.intermediate_size, + self.hidden_size // 32, + dtype=torch.uint8, + ), + ) + self.register_buffer( + "w2_scale_inv", + torch.empty( + self.num_local_experts, + self.hidden_size, + self.intermediate_size // 32, + dtype=torch.uint8, + ), + ) + self._loaded: set[tuple[int, str]] = set() + self._prepared = False + + def is_local_expert(self, expert_id: int) -> bool: + return self.local_expert_start <= int(expert_id) < self.local_expert_end + + def load_expert_weight( + self, + expert_id: int, + projection: str, + weight: torch.Tensor, + scale: torch.Tensor | None, + ) -> None: + expert_id = int(expert_id) + if not self.is_local_expert(expert_id): + raise ValueError(f"Expert {expert_id} is not owned by this EP rank.") + key = (expert_id, projection) + if key in self._loaded: + raise ValueError(f"Duplicate DeepSeek V4 expert tensor {key}.") + self.provider.load_expert_projection( + self.op_spec, + local_expert_id=expert_id - self.local_expert_start, + projection=projection, + loaded_weight=weight, + loaded_scale=scale, + w13_weight=self.w13_weight.data, + w2_weight=self.w2_weight.data, + w13_scale_inv=self.w13_scale_inv, + w2_scale_inv=self.w2_scale_inv, + ) + self._loaded.add(key) + + def validate_loaded_weights(self) -> None: + expected = { + (expert_id, projection) + for expert_id in range(self.local_expert_start, self.local_expert_end) + for projection in ("gate", "down", "up") + } + missing = sorted(expected - self._loaded) + if missing: + raise ValueError(f"Missing local DeepSeek V4 expert tensors: {missing[:8]}.") + + def prepare_for_inference(self) -> None: + if self._prepared: + return + w13, w2, s13, s2 = self.provider.prepare_weights( + self.w13_weight, + self.w2_weight, + self.w13_scale_inv, + self.w2_scale_inv, + ) + self.w13_weight.data = w13 + self.w2_weight.data = w2 + self.w13_scale_inv = s13 + self.w2_scale_inv = s2 + self._prepared = True + + def forward( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> torch.Tensor: + if not self._prepared: + raise RuntimeError("DeepSeek V4 FP4 experts were not prepared after loading.") + return self.provider.run( + self.op_spec, + hidden_states, + topk_ids, + topk_weights, + self.w13_weight, + self.w2_weight, + self.w13_scale_inv, + self.w2_scale_inv, + local_expert_start=self.local_expert_start, + ep_rank=self.ep_rank, + ) + + +class DeepseekV4Moe(nn.Module): + def __init__(self, config, layer_idx: int) -> None: + super().__init__() + self.parallel = get_parallel_context() + self.gate = DeepseekV4Router( + config, hash_routing=config.mlp_layer_types[int(layer_idx)] == "hash_moe" + ) + self.experts = DeepseekV4PackedExperts(config) + self.shared_experts = DeepseekV4SharedExperts(config) + + def forward(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: + output_shape = hidden_states.shape + hidden_states = hidden_states.reshape(-1, output_shape[-1]) + input_ids = input_ids.reshape(-1) + shared = self.shared_experts(hidden_states) + global_hidden = self.parallel.ep_all_gather_into_tensor(hidden_states) + global_input_ids = self.parallel.ep_all_gather_into_tensor(input_ids) + topk_weights, topk_ids = self.gate(global_hidden, global_input_ids) + routed = self.experts(global_hidden, topk_ids, topk_weights) + routed = self.parallel.ep_reduce_scatter_tensor(routed) + return (routed + shared).view(output_shape) + + +class DeepseekV4Attention(ReferenceDeepseekV4Attention): + """Reference DSV4 attention math backed by Sparse-vLLM FP8 operators.""" + + def __init__(self, config, layer_idx: int) -> None: + super().__init__(config, layer_idx) + quantization = config.quantization_config + self.q_a_proj = ReplicatedLinear( + int(config.hidden_size), int(config.q_lora_rank), quantization=quantization + ) + self.q_b_proj = ReplicatedLinear( + int(config.q_lora_rank), + int(config.num_attention_heads * config.head_dim), + quantization=quantization, + ) + self.kv_proj = ReplicatedLinear( + int(config.hidden_size), int(config.head_dim), quantization=quantization + ) + self.o_a_proj = DeepseekV4GroupedFp8Linear(config) + self.o_b_proj = ReplicatedLinear( + int(config.o_groups * config.o_lora_rank), + int(config.hidden_size), + quantization=quantization, + ) + if self.compressor is not None and hasattr(self.compressor, "indexer"): + self.compressor.indexer.q_b_proj = ReplicatedLinear( + int(config.q_lora_rank), + int(config.index_n_heads * config.index_head_dim), + quantization=quantization, + ) + + @staticmethod + def _store_rows(cache: torch.Tensor, rows: torch.Tensor, columns: torch.Tensor, values: torch.Tensor) -> None: + cache[rows.long(), columns.long()] = values + + def _compress_step( + self, + module, + hidden_states: torch.Tensor, + positions: torch.Tensor, + rows: torch.Tensor, + *, + indexer: bool = False, + ) -> None: + manager = get_context().cache_manager + layer_idx = int(self.layer_idx) + is_csa = self.layer_type == "compressed_sparse_attention" + ratio = int(module.compress_rate) + head_dim = int(module.head_dim) + slot = manager.csa_slot(layer_idx) if is_csa else manager.hca_slot(layer_idx) + projected_kv = module.kv_proj(hidden_states) + projected_gate = module.gate_proj(hidden_states) + + if is_csa: + prefix = "index" if indexer else "csa" + ring_kv = getattr(manager, f"{prefix}_ring_kv")[slot] + ring_gate = getattr(manager, f"{prefix}_ring_gate")[slot] + overlap_kv = getattr(manager, f"{prefix}_overlap_kv")[slot] + overlap_gate = getattr(manager, f"{prefix}_overlap_gate")[slot] + output_cache = manager.csa_index[slot] if indexer else manager.csa_kv[slot] + else: + ring_kv = manager.hca_ring_kv[slot] + ring_gate = manager.hca_ring_gate[slot] + output_cache = manager.hca_kv[slot] + + for token_idx in range(hidden_states.shape[1]): + position = positions[:, token_idx] + ring_column = torch.remainder(position, ratio).long() + kv_token = projected_kv[:, token_idx] + gate_token = projected_gate[:, token_idx] + module.position_bias[ring_column] + self._store_rows(ring_kv, rows, ring_column, kv_token) + self._store_rows(ring_gate, rows, ring_column, gate_token) + current_kv = ring_kv[rows.long()] + current_gate = ring_gate[rows.long()] + + if is_csa: + combined_kv = torch.cat( + [overlap_kv[rows.long()], current_kv[..., head_dim:]], dim=1 + ) + combined_gate = torch.cat( + [overlap_gate[rows.long()], current_gate[..., head_dim:]], dim=1 + ) + next_overlap_kv = current_kv[..., :head_dim] + next_overlap_gate = current_gate[..., :head_dim] + else: + combined_kv = current_kv + combined_gate = current_gate + + compressed = module.kv_norm( + ( + combined_kv + * combined_gate.softmax(dim=1, dtype=torch.float32).to(combined_kv.dtype) + ).sum(dim=1) + ) + window_position = (position - ratio + 1).clamp_min(0).unsqueeze(1) + cos, sin = module.rotary_emb( + compressed.unsqueeze(1), + position_ids=window_position, + layer_type=module.rope_layer_type, + ) + compressed = apply_rotary_pos_emb( + compressed[:, None, None, :], cos, sin + )[:, 0, 0] + closes_window = torch.remainder(position + 1, ratio).eq(0) + entry = torch.div(position, ratio, rounding_mode="floor").clamp( + max=output_cache.shape[1] - 1 + ) + old = output_cache[rows.long(), entry.long()] + self._store_rows( + output_cache, + rows, + entry, + torch.where(closes_window[:, None], compressed, old), + ) + if is_csa: + old_kv = overlap_kv[rows.long()] + old_gate = overlap_gate[rows.long()] + overlap_kv[rows.long()] = torch.where( + closes_window[:, None, None], next_overlap_kv, old_kv + ) + overlap_gate[rows.long()] = torch.where( + closes_window[:, None, None], next_overlap_gate, old_gate + ) + + def _compressed_kv( + self, + hidden_states: torch.Tensor, + q_residual: torch.Tensor, + positions: torch.Tensor, + rows: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + manager = get_context().cache_manager + batch, seq_len, _ = hidden_states.shape + if self.compressor is None: + return hidden_states.new_empty((batch, seq_len, 0, self.head_dim)), torch.empty( + (batch, seq_len, 0), dtype=torch.bool, device=hidden_states.device + ) + + self._compress_step(self.compressor, hidden_states, positions, rows) + ratio = int(self.compressor.compress_rate) + capacity = manager.compressed_capacity(ratio, positions) + entry_ids = torch.arange(capacity, device=hidden_states.device) + threshold = torch.div(positions + 1, ratio, rounding_mode="floor") + + if self.layer_type == "heavily_compressed_attention": + cache = manager.hca_kv[manager.hca_slot(self.layer_idx), rows.long(), :capacity] + valid = entry_ids.view(1, 1, -1) < threshold.unsqueeze(-1) + return cache.unsqueeze(1).expand(-1, seq_len, -1, -1), valid + + indexer = self.compressor.indexer + self._compress_step(indexer, hidden_states, positions, rows, indexer=True) + slot = manager.csa_slot(self.layer_idx) + index_cache = manager.csa_index[slot, rows.long(), :capacity] + outer_cache = manager.csa_kv[slot, rows.long(), :capacity] + if capacity == 0: + return outer_cache.unsqueeze(1).expand(-1, seq_len, -1, -1), torch.empty( + (batch, seq_len, 0), dtype=torch.bool, device=hidden_states.device + ) + + cos, sin = indexer.rotary_emb( + hidden_states, position_ids=positions, layer_type=indexer.rope_layer_type + ) + index_q = indexer.q_b_proj(q_residual).view( + batch, seq_len, indexer.num_heads, indexer.head_dim + ) + index_q = apply_rotary_pos_emb(index_q.transpose(1, 2), cos, sin).transpose(1, 2) + scores = torch.einsum("bshd,btd->bsht", index_q.float(), index_cache.float()) + scores = F.relu(scores) * indexer.scorer.softmax_scale + weights = indexer.scorer.weights_proj(hidden_states).float() * indexer.scorer.weights_scaling + scores = (scores * weights.unsqueeze(-1)).sum(dim=2) + valid_entries = entry_ids.view(1, 1, -1) < threshold.unsqueeze(-1) + scores = scores.masked_fill(~valid_entries, float("-inf")) + top_k = min(int(indexer.index_topk), capacity) + indices = scores.topk(top_k, dim=-1).indices + valid = indices < threshold.unsqueeze(-1) + gather_index = indices.unsqueeze(-1).expand(-1, -1, -1, self.head_dim) + selected = torch.gather( + outer_cache.unsqueeze(1).expand(-1, seq_len, -1, -1), 2, gather_index + ) + return selected, valid + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings, + position_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values=None, + **kwargs, + ) -> tuple[torch.Tensor, None]: + del attention_mask, past_key_values + context = get_context() + manager = context.cache_manager + state = manager.get_layer_batch_states(self.layer_idx) + rows = kwargs.pop("cache_rows", state.req_indices).long() + if kwargs: + raise TypeError(f"Unexpected DeepSeek V4 attention arguments: {sorted(kwargs)}") + batch, seq_len, _ = hidden_states.shape + if rows.numel() != batch: + raise RuntimeError( + f"DeepSeek V4 attention row count mismatch: rows={rows.numel()} batch={batch}." + ) + + cos, sin = position_embeddings[self.rope_layer_type] + q_residual = self.q_a_norm(self.q_a_proj(hidden_states)) + q = self.q_b_proj(q_residual).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + q = apply_rotary_pos_emb(self.q_b_norm(q), cos, sin) + kv = self.kv_proj(hidden_states).view(batch, seq_len, 1, self.head_dim).transpose(1, 2) + kv = apply_rotary_pos_emb(self.kv_norm(kv), cos, sin)[:, 0] + + prefix_len = self.sliding_window - 1 + prefix_offsets = torch.arange( + prefix_len, 0, -1, device=hidden_states.device + ) + prefix_positions = position_ids[:, :1] - prefix_offsets + prefix_valid = prefix_positions >= 0 + prefix_columns = torch.remainder(prefix_positions, self.sliding_window).long() + raw_layer = manager.raw_kv[self.layer_idx] + prefix = raw_layer[rows[:, None], prefix_columns] + raw = torch.cat([prefix, kv], dim=1) + raw_positions = torch.cat([prefix_positions, position_ids], dim=1) + raw_valid = raw_positions.unsqueeze(1) <= position_ids.unsqueeze(-1) + raw_valid &= raw_positions.unsqueeze(1) >= position_ids.unsqueeze(-1) - self.sliding_window + 1 + raw_valid[:, :, :prefix_len] &= prefix_valid.unsqueeze(1) + + compressed, compressed_valid = self._compressed_kv( + hidden_states, q_residual, position_ids, rows + ) + raw_per_query = raw.unsqueeze(1).expand(-1, seq_len, -1, -1) + all_kv = torch.cat([raw_per_query, compressed], dim=2) + valid = torch.cat([raw_valid, compressed_valid], dim=-1) + scores = torch.einsum("bhsd,bskd->bhsk", q, all_kv) * self.scaling + scores = scores.masked_fill(~valid.unsqueeze(1), float("-inf")) + sinks = self.sinks.view(1, -1, 1, 1).expand(batch, -1, seq_len, -1) + logits = torch.cat([scores, sinks], dim=-1) + logits = logits - logits.max(dim=-1, keepdim=True).values + probs = F.softmax(logits, dim=-1)[..., :-1].to(all_kv.dtype) + output = torch.einsum("bhsk,bskd->bhsd", probs, all_kv) + output = apply_rotary_pos_emb(output, cos, -sin).transpose(1, 2) + + columns = torch.remainder(position_ids, self.sliding_window).long() + # A long prefill revisits ring columns. Advanced assignment with + # duplicate indices has undefined winner ordering on both CPU and CUDA, + # so persist only the final, column-unique window. + store_start = max(0, seq_len - self.sliding_window) + raw_layer[rows[:, None], columns[:, store_start:]] = kv[:, store_start:] + grouped = output.reshape(batch, seq_len, self.config.o_groups, -1) + grouped = self.o_a_proj(grouped).flatten(2) + return self.o_b_proj(grouped), None + + +class DeepseekV4DecoderLayer(nn.Module): + def __init__(self, config, layer_idx: int) -> None: + super().__init__() + self.attn = DeepseekV4Attention(config, layer_idx) + self.ffn = DeepseekV4Moe(config, layer_idx) + self.attn_norm = DeepseekV4RMSNorm( + int(config.hidden_size), eps=float(config.rms_norm_eps) + ) + self.ffn_norm = DeepseekV4RMSNorm( + int(config.hidden_size), eps=float(config.rms_norm_eps) + ) + self.hc_attn = DeepseekV4HyperConnection(config) + self.hc_ffn = DeepseekV4HyperConnection(config) + + def forward( + self, + hidden_states: torch.Tensor, + *, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + position_embeddings, + attention_mask: torch.Tensor, + past_key_values, + cache_rows: torch.Tensor | None = None, + ) -> torch.Tensor: + dtype = hidden_states.dtype + post, comb, collapsed = self.hc_attn(hidden_states) + attention_output, _ = self.attn( + self.attn_norm(collapsed), + position_embeddings=position_embeddings, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + cache_rows=cache_rows, + ) + hidden_states = post.to(dtype).unsqueeze(-1) * attention_output.unsqueeze( + -2 + ) + torch.matmul(comb.to(dtype).transpose(-1, -2), hidden_states) + post, comb, collapsed = self.hc_ffn(hidden_states) + ffn_output = self.ffn(self.ffn_norm(collapsed), input_ids) + return post.to(dtype).unsqueeze(-1) * ffn_output.unsqueeze(-2) + torch.matmul( + comb.to(dtype).transpose(-1, -2), hidden_states + ) + + +class DeepseekV4Model(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.config = config + self.embed = nn.Embedding(int(config.vocab_size), int(config.hidden_size)) + self.layers = nn.ModuleList( + DeepseekV4DecoderLayer(config, layer_idx) + for layer_idx in range(int(config.num_hidden_layers)) + ) + self.hc_head = DeepseekV4HyperHead(config) + self.norm = DeepseekV4RMSNorm( + int(config.hidden_size), eps=float(config.rms_norm_eps) + ) + self.rotary_emb = DeepseekV4RotaryEmbedding(config) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + cache_rows: torch.Tensor, + ) -> torch.Tensor: + if input_ids.ndim != 2 or position_ids.shape != input_ids.shape: + raise ValueError( + "Native DeepSeek V4 expects [batch, sequence] input and position tensors, " + f"got {tuple(input_ids.shape)} and {tuple(position_ids.shape)}." + ) + inputs_embeds = self.embed(input_ids) + position_embeddings = { + "main": self.rotary_emb( + inputs_embeds, position_ids=position_ids, layer_type="main" + ), + "compress": self.rotary_emb( + inputs_embeds, position_ids=position_ids, layer_type="compress" + ), + } + hidden_states = inputs_embeds.unsqueeze(2).expand( + -1, -1, int(self.config.hc_mult), -1 + ).contiguous() + for layer in self.layers: + hidden_states = layer( + hidden_states, + input_ids=input_ids, + position_ids=position_ids, + position_embeddings=position_embeddings, + attention_mask=None, + past_key_values=None, + cache_rows=cache_rows, + ) + return self.norm(self.hc_head(hidden_states)) diff --git a/src/sparsevllm/models/layout.py b/src/sparsevllm/models/layout.py new file mode 100644 index 00000000..91890743 --- /dev/null +++ b/src/sparsevllm/models/layout.py @@ -0,0 +1,221 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sparsevllm.utils.config import config_get + + +def _coerce_int_list( + name: str, + value: Any, + *, + allow_none: bool = False, +) -> list[int] | None: + if value is None: + if allow_none: + return None + raise ValueError(f"{name} is required.") + if isinstance(value, str): + return [int(part) for part in value.split(",") if part.strip()] + if isinstance(value, (list, tuple)): + return [int(item) for item in value] + raise ValueError( + f"{name} must be a list/tuple of ints or a comma-separated string, " + f"got {value!r}." + ) + + +def _attention_type(value: Any) -> str: + value = str(value).strip().lower() + if value in { + "full", + "full_attention", + "attention", + "self_attention", + "sliding_attention", + "compressed_sparse_attention", + "heavily_compressed_attention", + }: + return "full" + if value in { + "linear", + "linear_attention", + "recurrent", + "recurrent_attention", + "gated_delta", + "gated_delta_net", + }: + return "linear" + raise ValueError(f"Unsupported attention layer type {value!r}.") + + +@dataclass(frozen=True) +class RuntimeLayout: + num_layers: int + num_kv_layers: int + full_attention_layer_indices: tuple[int, ...] + linear_attention_layer_indices: tuple[int, ...] + layer_idx_to_kv_idx: tuple[int | None, ...] + kv_idx_to_layer_idx: tuple[int, ...] + + @classmethod + def dense(cls, num_layers: int) -> "RuntimeLayout": + num_layers = int(num_layers) + if num_layers <= 0: + raise ValueError(f"num_hidden_layers must be positive, got {num_layers}.") + layers = tuple(range(num_layers)) + return cls( + num_layers=num_layers, + num_kv_layers=num_layers, + full_attention_layer_indices=layers, + linear_attention_layer_indices=(), + layer_idx_to_kv_idx=tuple(range(num_layers)), + kv_idx_to_layer_idx=layers, + ) + + @classmethod + def from_config( + cls, + hf_config: Any, + *, + require_mixed: bool = False, + ) -> "RuntimeLayout": + num_layers = int(config_get(hf_config, "num_hidden_layers")) + if num_layers <= 0: + raise ValueError(f"num_hidden_layers must be positive, got {num_layers}.") + layer_types = config_get(hf_config, "layer_types", None) + full_layers = _coerce_int_list( + "full_attention_layer_indices", + config_get( + hf_config, + "full_attention_layer_indices", + config_get(hf_config, "attention_layer_indices", None), + ), + allow_none=True, + ) + linear_layers = _coerce_int_list( + "linear_attention_layer_indices", + config_get(hf_config, "linear_attention_layer_indices", None), + allow_none=True, + ) + + if layer_types is not None: + if len(layer_types) != num_layers: + raise ValueError( + "layer_types length must equal num_hidden_layers: " + f"{len(layer_types)} != {num_layers}." + ) + inferred_full, inferred_linear = [], [] + for layer_idx, layer_type in enumerate(layer_types): + target = ( + inferred_full + if _attention_type(layer_type) == "full" + else inferred_linear + ) + target.append(layer_idx) + full_layers = inferred_full if full_layers is None else full_layers + linear_layers = inferred_linear if linear_layers is None else linear_layers + + if full_layers is None and linear_layers is None: + if require_mixed: + raise ValueError( + "Mixed-attention models require layer_types or explicit " + "full/linear attention layer indices." + ) + return cls.dense(num_layers) + if full_layers is None: + linear_set = set(linear_layers or ()) + full_layers = [idx for idx in range(num_layers) if idx not in linear_set] + if linear_layers is None: + full_set = set(full_layers or ()) + linear_layers = [idx for idx in range(num_layers) if idx not in full_set] + + full_tuple = tuple(sorted(int(idx) for idx in full_layers)) + linear_tuple = tuple(sorted(int(idx) for idx in linear_layers)) + full_set, linear_set = set(full_tuple), set(linear_tuple) + expected = set(range(num_layers)) + if full_set & linear_set: + raise ValueError( + "RuntimeLayout full and linear layer sets overlap: " + f"{sorted(full_set & linear_set)}." + ) + if full_set | linear_set != expected: + raise ValueError( + "RuntimeLayout layer map is incomplete: " + f"missing={sorted(expected - (full_set | linear_set))}, " + f"extra={sorted((full_set | linear_set) - expected)}." + ) + + raw_layer_to_kv = config_get(hf_config, "layer_idx_to_kv_idx", None) + if raw_layer_to_kv is None: + layer_to_kv: list[int | None] = [None] * num_layers + for kv_idx, layer_idx in enumerate(full_tuple): + layer_to_kv[layer_idx] = kv_idx + else: + if len(raw_layer_to_kv) != num_layers: + raise ValueError( + "layer_idx_to_kv_idx length must equal num_hidden_layers: " + f"{len(raw_layer_to_kv)} != {num_layers}." + ) + layer_to_kv = [ + None if value is None or int(value) < 0 else int(value) + for value in raw_layer_to_kv + ] + invalid_linear = [ + layer_idx + for layer_idx in linear_tuple + if layer_to_kv[layer_idx] is not None + ] + if invalid_linear: + raise ValueError( + "Linear-attention layers must not have KV indices: " + f"{invalid_linear}." + ) + + kv_pairs = sorted( + (kv_idx, layer_idx) + for layer_idx, kv_idx in enumerate(layer_to_kv) + if kv_idx is not None + ) + if len(kv_pairs) != len(full_tuple): + raise ValueError( + "RuntimeLayout must assign one KV index to each full-attention " + f"layer: full={len(full_tuple)}, assigned={len(kv_pairs)}." + ) + kv_indices = [kv_idx for kv_idx, _ in kv_pairs] + if kv_indices != list(range(len(kv_pairs))): + raise ValueError(f"KV layer indices must be contiguous, got {kv_indices}.") + kv_tuple = tuple(layer_idx for _, layer_idx in kv_pairs) + configured_num_kv_layers = config_get(hf_config, "num_kv_layers", None) + if ( + configured_num_kv_layers is not None + and int(configured_num_kv_layers) != len(kv_tuple) + ): + raise ValueError( + f"num_kv_layers={configured_num_kv_layers} does not match " + f"full-attention layers={len(kv_tuple)}." + ) + return cls( + num_layers=num_layers, + num_kv_layers=len(kv_tuple), + full_attention_layer_indices=full_tuple, + linear_attention_layer_indices=linear_tuple, + layer_idx_to_kv_idx=tuple(layer_to_kv), + kv_idx_to_layer_idx=kv_tuple, + ) + + def is_full_attention(self, layer_idx: int) -> bool: + return self.layer_idx_to_kv_idx[int(layer_idx)] is not None + + def is_linear_attention(self, layer_idx: int) -> bool: + return self.layer_idx_to_kv_idx[int(layer_idx)] is None + + def kv_layer_index(self, layer_idx: int) -> int: + layer_idx = int(layer_idx) + kv_idx = self.layer_idx_to_kv_idx[layer_idx] + if kv_idx is None: + raise RuntimeError( + f"layer_idx={layer_idx} is linear_attention and has no KV cache." + ) + return int(kv_idx) diff --git a/src/sparsevllm/models/qwen3_5.py b/src/sparsevllm/models/qwen3_5.py index 97e1c218..2ba37904 100644 --- a/src/sparsevllm/models/qwen3_5.py +++ b/src/sparsevllm/models/qwen3_5.py @@ -7,7 +7,6 @@ import torch.nn.functional as F from sparsevllm.distributed import get_parallel_context -from sparsevllm.layers.activation import SiluAndMul from sparsevllm.layers.attention import Attention from sparsevllm.layers.layernorm import GemmaRMSNorm from sparsevllm.layers.linear import ( @@ -16,6 +15,10 @@ RowParallelLinear, divide, ) +from sparsevllm.operators.gate_up_swiglu import ( + GateUpSwiGLUOpSpec, + resolve_gate_up_swiglu_provider, +) from sparsevllm.layers.rotary_embedding import apply_partial_rotary_emb, get_rope from sparsevllm.layers.embed_head import VocabParallelEmbedding, ParallelLMHead from sparsevllm.utils.context import get_context @@ -375,42 +378,28 @@ def __init__(self, config) -> None: if self.activation not in ("silu", "swish"): raise NotImplementedError(f"qwen3_5 linear attention supports silu/swish activation, got {self.activation!r}.") self.proj_chunk_size = int(getattr(config, "mlp_chunk_size", 16384)) + self.recurrent_state_dtype = getattr( + config, + "runtime_recurrent_state_dtype", + getattr(config, "torch_dtype", torch.bfloat16), + ) + if not isinstance(self.recurrent_state_dtype, torch.dtype): + raise TypeError( + "qwen3_5 runtime_recurrent_state_dtype must be a torch.dtype, " + f"got {self.recurrent_state_dtype!r}." + ) quantization = getattr(config, "quantization_config", None) hidden_size = int(config.hidden_size) - self.in_proj_q = ColumnParallelLinear( - hidden_size, - self.key_dim, - bias=False, - quantization=quantization, - ) - self.in_proj_k = ColumnParallelLinear( - hidden_size, - self.key_dim, - bias=False, - quantization=quantization, - ) - self.in_proj_v = ColumnParallelLinear( - hidden_size, - self.value_dim, - bias=False, - quantization=quantization, - ) - self.in_proj_z = ColumnParallelLinear( + self.in_proj_qkvz = MergedColumnParallelLinear( hidden_size, - self.value_dim, + [self.key_dim, self.key_dim, self.value_dim, self.value_dim], bias=False, quantization=quantization, ) - self.in_proj_b = ColumnParallelLinear( - hidden_size, - self.total_num_v_heads, - bias=False, - quantization=None, - ) - self.in_proj_a = ColumnParallelLinear( + self.in_proj_ba = MergedColumnParallelLinear( hidden_size, - self.total_num_v_heads, + [self.total_num_v_heads, self.total_num_v_heads], bias=False, quantization=None, ) @@ -462,16 +451,21 @@ def _split_row_block_scale(loaded_scale: torch.Tensor | None, row_sizes: list[in def _load_projection_weight( self, - module: ColumnParallelLinear, + module: MergedColumnParallelLinear, loaded_weight: torch.Tensor, loaded_scale: torch.Tensor | None, + shard_id: int, ) -> None: if loaded_scale is not None: - module.load_quantized_weight(loaded_weight, loaded_scale) + module.load_quantized_weight( + loaded_weight, + loaded_scale, + loaded_shard_id=shard_id, + ) return if bool(getattr(module, "quantized", False)): raise ValueError(f"Missing FP8 weight_scale_inv for quantized qwen3_5 projection {module}.") - module.weight_loader(module.weight, loaded_weight) + module.weight_loader(module.weight, loaded_weight, shard_id) def load_packed_in_proj_qkv( self, @@ -485,9 +479,15 @@ def load_packed_in_proj_qkv( ) q, k, v = torch.split(loaded_weight, row_sizes, dim=0) q_scale, k_scale, v_scale = self._split_row_block_scale(loaded_scale, row_sizes) - self._load_projection_weight(self.in_proj_q, q, q_scale) - self._load_projection_weight(self.in_proj_k, k, k_scale) - self._load_projection_weight(self.in_proj_v, v, v_scale) + for shard_id, (weight, scale) in enumerate( + ((q, q_scale), (k, k_scale), (v, v_scale)) + ): + self._load_projection_weight( + self.in_proj_qkvz, + weight, + scale, + shard_id, + ) return 3 def _split_interleaved_qkvz( @@ -546,10 +546,15 @@ def load_packed_in_proj_qkvz( ) -> int: q, k, v, z = self._split_interleaved_qkvz(loaded_weight) q_scale, k_scale, v_scale, z_scale = self._split_interleaved_qkvz_scale(loaded_scale) - self._load_projection_weight(self.in_proj_q, q, q_scale) - self._load_projection_weight(self.in_proj_k, k, k_scale) - self._load_projection_weight(self.in_proj_v, v, v_scale) - self._load_projection_weight(self.in_proj_z, z, z_scale) + for shard_id, (weight, scale) in enumerate( + ((q, q_scale), (k, k_scale), (v, v_scale), (z, z_scale)) + ): + self._load_projection_weight( + self.in_proj_qkvz, + weight, + scale, + shard_id, + ) return 4 def load_packed_in_proj_ba( @@ -569,8 +574,8 @@ def load_packed_in_proj_ba( b = weight[:, :num_v_per_k, :].reshape(-1, hidden) a = weight[:, num_v_per_k:, :].reshape(-1, hidden) b_scale, a_scale = self._split_row_block_scale(loaded_scale, [self.total_num_v_heads, self.total_num_v_heads]) - self._load_projection_weight(self.in_proj_b, b, b_scale) - self._load_projection_weight(self.in_proj_a, a, a_scale) + self._load_projection_weight(self.in_proj_ba, b, b_scale, 0) + self._load_projection_weight(self.in_proj_ba, a, a_scale, 1) return 2 def _tp_vector_weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor) -> None: @@ -585,13 +590,11 @@ def _tp_vector_weight_loader(self, param: nn.Parameter, loaded_weight: torch.Ten param.data.copy_(loaded_weight.reshape(-1).narrow(0, start, shard_size).to(dtype=param.dtype)) def _project_qkvzba(self, hidden_states: torch.Tensor): - q = self.in_proj_q(hidden_states) - k = self.in_proj_k(hidden_states) - v = self.in_proj_v(hidden_states) - z = self.in_proj_z(hidden_states) - b = self.in_proj_b(hidden_states) - a = self.in_proj_a(hidden_states) - mixed_qkv = torch.cat([q, k, v], dim=-1) + mixed_qkv, z = self.in_proj_qkvz(hidden_states).split( + [2 * self.tp_key_dim + self.tp_value_dim, self.tp_value_dim], + dim=-1, + ) + b, a = self.in_proj_ba(hidden_states).chunk(2, dim=-1) z = z.view(-1, self.num_v_heads, self.head_v_dim) return mixed_qkv, z, b, a @@ -625,7 +628,7 @@ def _load_batch_states( if seqs is None: raise RuntimeError("qwen3_5 linear attention requires context.seqs.") conv_dtype = activation_dtype - recurrent_dtype = activation_dtype + recurrent_dtype = self.recurrent_state_dtype conv_states = [] recurrent_states = [] has_initial = [] @@ -663,7 +666,9 @@ def _store_batch_states(self, context, recurrent_state_manager, conv_states: tor layer_idx, { "conv_state": conv_states[row], - "recurrent_state": recurrent_states[row].to(dtype=conv_states.dtype), + "recurrent_state": recurrent_states[row].to( + dtype=self.recurrent_state_dtype + ), }, ) @@ -811,28 +816,59 @@ def forward(self, positions: torch.Tensor, hidden_states: torch.Tensor) -> torch class Qwen35MLP(nn.Module): - def __init__(self, config) -> None: + def __init__( + self, + config, + *, + intermediate_size: int | None = None, + reduce_results: bool = True, + ) -> None: super().__init__() + intermediate_size = int( + config.intermediate_size if intermediate_size is None else intermediate_size + ) quantization = getattr(config, "quantization_config", None) self.gate_up_proj = MergedColumnParallelLinear( int(config.hidden_size), - [int(config.intermediate_size)] * 2, + [intermediate_size] * 2, bias=False, quantization=quantization, ) self.down_proj = RowParallelLinear( - int(config.intermediate_size), + intermediate_size, int(config.hidden_size), bias=False, quantization=quantization, + reduce_results=reduce_results, ) if getattr(config, "hidden_act", "silu") != "silu": raise NotImplementedError(f"qwen3_5 supports hidden_act='silu', got {config.hidden_act!r}.") - self.act_fn = SiluAndMul() + activation_dtype = getattr(config, "dtype", None) or getattr( + config, "torch_dtype", torch.bfloat16 + ) + self.gate_up_swiglu_spec = GateUpSwiGLUOpSpec( + hidden_size=int(config.hidden_size), + intermediate_size=intermediate_size, + tp_size=self.gate_up_proj.tp_size, + activation_dtype=activation_dtype, + weight_dtype=( + torch.float8_e4m3fn + if self.gate_up_proj.quantized + else activation_dtype + ), + cuda_graph=bool(getattr(config, "decode_cuda_graph", False)), + ) + self.gate_up_swiglu_provider = resolve_gate_up_swiglu_provider( + self.gate_up_swiglu_spec + ) self.mlp_chunk_size = int(getattr(config, "mlp_chunk_size", 16384)) def _forward_chunk(self, x: torch.Tensor) -> torch.Tensor: - return self.down_proj(self.act_fn(self.gate_up_proj(x))) + return self.down_proj( + self.gate_up_swiglu_provider.run( + self.gate_up_swiglu_spec, x, self.gate_up_proj + ) + ) def forward(self, x: torch.Tensor) -> torch.Tensor: if int(x.shape[0]) <= self.mlp_chunk_size: @@ -845,7 +881,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class Qwen35DecoderLayer(nn.Module): - def __init__(self, config, layer_idx: int) -> None: + def __init__(self, config, layer_idx: int, mlp_cls=Qwen35MLP) -> None: super().__init__() self.layer_idx = int(layer_idx) runtime_layout = getattr(config, "runtime_layout", None) @@ -857,7 +893,7 @@ def __init__(self, config, layer_idx: int) -> None: else: self.linear_attn = Qwen35LinearAttention(config) self.attention_type = "linear_attention" - self.mlp = Qwen35MLP(config) + self.mlp = mlp_cls(config) self.input_layernorm = Qwen35RMSNorm(int(config.hidden_size), eps=float(config.rms_norm_eps)) self.post_attention_layernorm = Qwen35RMSNorm(int(config.hidden_size), eps=float(config.rms_norm_eps)) @@ -881,12 +917,12 @@ def forward( class Qwen35Model(nn.Module): - def __init__(self, config) -> None: + def __init__(self, config, layer_cls=Qwen35DecoderLayer) -> None: super().__init__() self.config = config self.embed_tokens = VocabParallelEmbedding(int(config.vocab_size), int(config.hidden_size)) self.layers = nn.ModuleList( - [Qwen35DecoderLayer(config, layer_idx) for layer_idx in range(int(config.num_hidden_layers))] + [layer_cls(config, layer_idx) for layer_idx in range(int(config.num_hidden_layers))] ) self.norm = Qwen35RMSNorm(int(config.hidden_size), eps=float(config.rms_norm_eps)) self.sparse_controller = None @@ -936,6 +972,10 @@ class Qwen35ForCausalLM(nn.Module): "q_proj": ("qkv_gate_proj", "q"), "k_proj": ("qkv_gate_proj", "k"), "v_proj": ("qkv_gate_proj", "v"), + "in_proj_z": ("in_proj_qkvz", 3), + "in_proj_b": ("in_proj_ba", 0), + "in_proj_a": ("in_proj_ba", 1), + "shared_expert_gate": ("gate", "shared"), "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), } diff --git a/src/sparsevllm/models/qwen3_5_moe.py b/src/sparsevllm/models/qwen3_5_moe.py new file mode 100644 index 00000000..f0edae0e --- /dev/null +++ b/src/sparsevllm/models/qwen3_5_moe.py @@ -0,0 +1,676 @@ +from __future__ import annotations + +import re + +import torch +import torch.nn.functional as F +from torch import nn + +from sparsevllm.distributed import get_parallel_context +from sparsevllm.engine.recurrent_state_manager import ( + RecurrentStateSpec, + RecurrentTensorSpec, +) +from sparsevllm.layers.embed_head import ParallelLMHead +from sparsevllm.models.qwen3_5 import ( + Qwen35DecoderLayer, + Qwen35ForCausalLM, + Qwen35MLP, + Qwen35Model, +) +from sparsevllm.models.qwen3_moe import Qwen3MoePackedExperts +from sparsevllm.operators.gated_shared_add import gated_shared_add +from sparsevllm.operators.moe import model_activation_dtype +from sparsevllm.operators.moe_router import ( + MoeRouterOpSpec, + resolve_moe_router_provider, +) +from sparsevllm.platforms import device_runtime +from sparsevllm.utils.log import logger +from sparsevllm.utils.weight_target import WeightTarget + + +_PACKED_EXPERT_SOURCE_RE = re.compile( + r"^model\.language_model\.layers\.(\d+)\.mlp\.experts\." + r"(gate_up_proj|down_proj)$" +) +_PACKED_EXPERT_TARGET_RE = re.compile( + r"^model\.layers\.(\d+)\.mlp\.experts\." + r"(gate_up_proj|down_proj)\.packed_expert_weight$" +) +_FP8_EXPERT_SOURCE_RE = re.compile( + r"^model\.language_model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." + r"(gate_proj|up_proj|down_proj)\.weight$" +) +_FP8_EXPERT_TARGET_RE = re.compile( + r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." + r"(gate_proj|up_proj|down_proj)\.expert_weight$" +) + + +class Qwen35MoeRouter(nn.Module): + """Replicated router with the checkpoint's FP32 softmax semantics.""" + + def __init__(self, config) -> None: + super().__init__() + self.hidden_size = int(config.hidden_size) + self.num_experts = int(config.num_experts) + self.top_k = int(config.num_experts_per_tok) + self.op_spec = MoeRouterOpSpec( + num_experts=self.num_experts, + top_k=self.top_k, + activation_dtype=model_activation_dtype(config), + norm_topk_prob=True, + cuda_graph=bool(getattr(config, "decode_cuda_graph", False)), + ) + self.provider = resolve_moe_router_provider(self.op_spec) + self.weight = nn.Parameter( + torch.empty(self.num_experts + 1, self.hidden_size) + ) + self.weight.weight_loader = self.weight_loader + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: str | None = None, + ) -> None: + target = param.data[-1:] if loaded_shard_id == "shared" else param.data[:-1] + if loaded_shard_id not in {None, "shared"} or target.shape != loaded_weight.shape: + raise ValueError( + "Qwen3.6 fused router/shared gate weight mismatch: " + f"shard={loaded_shard_id!r} expected={tuple(target.shape)} " + f"got={tuple(loaded_weight.shape)}." + ) + target.copy_(loaded_weight) + + def forward( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + router_logits, shared_gate_logits = F.linear( + hidden_states, + self.weight, + ).split((self.num_experts, 1), dim=-1) + topk_weights, topk_ids = self.provider.run( + self.op_spec, router_logits + ) + return topk_weights, topk_ids, shared_gate_logits + + +class Qwen35MoePackedExperts(Qwen3MoePackedExperts): + """Adapt Qwen3.6 BF16/FP8 checkpoints to the shared MoE provider.""" + + checkpoint_projection_map = { + "gate_up_proj": "gate_up", + "down_proj": "down", + "gate_proj": "gate", + "up_proj": "up", + } + + def __init__(self, config) -> None: + super().__init__(config) + self._loaded_packed_projections: set[str] = set() + + def rank_local_weight_slice( + self, + source_shape: tuple[int, ...], + *, + loaded_shard_id: str | tuple[int, str], + is_scale: bool = False, + ) -> tuple[slice, ...] | None: + if isinstance(loaded_shard_id, tuple): + if not self.fp8_enabled: + raise ValueError( + "Qwen3.6 BF16 checkpoints must use packed 3D expert weights." + ) + return super().rank_local_weight_slice( + source_shape, + loaded_shard_id=loaded_shard_id, + is_scale=is_scale, + ) + if self.fp8_enabled: + raise ValueError( + "Qwen3.6 FP8 checkpoints must use per-expert projections and scales." + ) + if is_scale: + raise ValueError("Qwen3.6 MoE BF16 experts do not have weight scales.") + expected = { + "gate_up_proj": ( + self.num_experts, + 2 * self.global_intermediate_size, + self.hidden_size, + ), + "down_proj": ( + self.num_experts, + self.hidden_size, + self.global_intermediate_size, + ), + }.get(str(loaded_shard_id)) + if expected is None: + raise ValueError( + f"Unsupported Qwen3.6 packed expert projection {loaded_shard_id!r}." + ) + if tuple(source_shape) != expected: + raise ValueError( + "Qwen3.6 packed expert checkpoint shape mismatch: " + f"projection={loaded_shard_id} expected={expected} " + f"got={tuple(source_shape)}." + ) + return ( + slice(self.local_expert_start, self.local_expert_end), + slice(None), + slice(None), + ) + + def load_packed_expert_weight( + self, + projection: str, + loaded_weight: torch.Tensor, + ) -> None: + if self.fp8_enabled: + raise ValueError( + "Qwen3.6 FP8 checkpoints must use per-expert projections and scales." + ) + projection = str(projection) + if projection in self._loaded_packed_projections: + raise ValueError( + f"Duplicate Qwen3.6 packed expert projection {projection!r}." + ) + if loaded_weight.dtype != torch.bfloat16: + raise TypeError( + "Qwen3.6 packed expert weights must be BF16, " + f"got {loaded_weight.dtype}." + ) + expected_local = { + "gate_up_proj": ( + self.num_local_experts, + 2 * self.global_intermediate_size, + self.hidden_size, + ), + "down_proj": ( + self.num_local_experts, + self.hidden_size, + self.global_intermediate_size, + ), + }.get(projection) + if expected_local is None: + raise ValueError( + f"Unsupported Qwen3.6 packed expert projection {projection!r}." + ) + if tuple(loaded_weight.shape) != expected_local: + raise ValueError( + "Qwen3.6 rank-local packed expert shape mismatch: " + f"projection={projection} expected={expected_local} " + f"got={tuple(loaded_weight.shape)}." + ) + + for local_expert_id in range(self.num_local_experts): + global_expert_id = self.local_expert_start + local_expert_id + if projection == "gate_up_proj": + gate, up = loaded_weight[local_expert_id].split( + self.global_intermediate_size, + dim=0, + ) + self.load_expert_weight( + global_expert_id, "gate_proj", gate, None + ) + self.load_expert_weight( + global_expert_id, "up_proj", up, None + ) + else: + self.load_expert_weight( + global_expert_id, + "down_proj", + loaded_weight[local_expert_id], + None, + ) + self._loaded_packed_projections.add(projection) + + def validate_loaded_weights(self) -> None: + if self.fp8_enabled: + super().validate_loaded_weights() + return + missing = {"gate_up_proj", "down_proj"} - self._loaded_packed_projections + if missing: + raise ValueError( + f"Missing Qwen3.6 packed expert projections: {sorted(missing)}." + ) + super().validate_loaded_weights() + + +class Qwen35MoeSparseMoeBlock(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.parallel_context = get_parallel_context() + self.mlp_chunk_size = int(getattr(config, "mlp_chunk_size", 16384)) + if self.mlp_chunk_size <= 0: + raise ValueError( + f"mlp_chunk_size must be > 0, got {self.mlp_chunk_size}." + ) + self.gate = Qwen35MoeRouter(config) + self.experts = Qwen35MoePackedExperts(config) + self.shared_expert = Qwen35MLP( + config, + intermediate_size=int(config.shared_expert_intermediate_size), + reduce_results=False, + ) + + def _forward_chunk( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + shared_output = self.shared_expert(hidden_states) + topk_weights, topk_ids, shared_gate_logits = self.gate(hidden_states) + local_output = self.experts(hidden_states, topk_ids, topk_weights) + if self.parallel_context.world.size > 1: + local_output, shared_output = self.parallel_context.world_all_reduce( + torch.stack((local_output, shared_output)) + ) + return gated_shared_add(local_output, shared_output, shared_gate_logits) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if hidden_states.dim() != 2: + raise ValueError( + "Qwen35MoeSparseMoeBlock expects [tokens, hidden], " + f"got {tuple(hidden_states.shape)}." + ) + chunks = hidden_states.split(self.mlp_chunk_size, dim=0) + outputs = [] + for chunk in chunks: + outputs.append(self._forward_chunk(chunk)) + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0) + + +class Qwen35MoeDecoderLayer(Qwen35DecoderLayer): + def __init__(self, config, layer_idx: int) -> None: + super().__init__(config, layer_idx, mlp_cls=Qwen35MoeSparseMoeBlock) + + +class Qwen35MoeModel(Qwen35Model): + def __init__(self, config) -> None: + setattr(config, "runtime_recurrent_state_dtype", torch.float32) + super().__init__(config, Qwen35MoeDecoderLayer) + + +class Qwen35MoeForCausalLM(Qwen35ForCausalLM): + ignored_weight_prefixes = ("model.visual.", "visual.", "mtp.") + special_weight_loaders = { + **Qwen35ForCausalLM.special_weight_loaders, + ".packed_expert_weight": "load_packed_expert_weight", + ".expert_weight": "load_expert_weight", + } + + def __init__(self, config) -> None: + nn.Module.__init__(self) + self.config = config + self.parallel_context = get_parallel_context() + self.model = Qwen35MoeModel(config) + self.lm_head = ParallelLMHead( + int(config.vocab_size), int(config.hidden_size) + ) + if bool(getattr(config, "tie_word_embeddings", False)): + self.lm_head.weight.data = self.model.embed_tokens.weight.data + self._loaded_linear_special_weights: set[str] = set() + self._intentionally_skipped_weights: set[str] = set() + self._intentionally_skipped_expert_weights: set[str] = set() + self._intentionally_skipped_expert_scales: set[str] = set() + + @staticmethod + def recurrent_state_spec(config, attention_tp_size: int) -> RecurrentStateSpec: + attention_tp_size = int(attention_tp_size) + num_k_heads = int(config.linear_num_key_heads) // attention_tp_size + num_v_heads = int(config.linear_num_value_heads) // attention_tp_size + key_head_dim = int(config.linear_key_head_dim) + value_head_dim = int(config.linear_value_head_dim) + conv_dim = 2 * num_k_heads * key_head_dim + num_v_heads * value_head_dim + return RecurrentStateSpec( + name="qwen3_5_moe gated delta net", + tensor_specs=( + RecurrentTensorSpec( + "conv_state", + (conv_dim, int(config.linear_conv_kernel_dim) - 1), + config.torch_dtype, + ), + RecurrentTensorSpec( + "recurrent_state", + (num_v_heads, key_head_dim, value_head_dim), + torch.float32, + ), + ), + ) + + @torch.inference_mode() + def warmup_moe(self, num_tokens: int = 1) -> None: + num_tokens = int(num_tokens) + if num_tokens <= 0: + raise ValueError(f"num_tokens must be > 0, got {num_tokens}.") + block = self.model.layers[0].mlp + experts = block.experts + device = experts.w13_weight.device + dtype = block.gate.weight.dtype + hidden_states = torch.zeros( + (num_tokens, experts.hidden_size), dtype=dtype, device=device + ) + top_k = int(self.config.num_experts_per_tok) + topk_ids = ( + torch.arange(num_tokens * top_k, dtype=torch.int64, device=device) + .remainder(experts.num_local_experts) + .add(experts.local_expert_start) + .view(num_tokens, top_k) + ) + topk_weights = torch.full( + (num_tokens, top_k), + 1.0 / top_k, + dtype=dtype, + device=device, + ) + experts(hidden_states, topk_ids, topk_weights) + block(hidden_states) + device_runtime.synchronize() + + def map_weight_name(self, source_weight_name: str) -> str | None: + match = _FP8_EXPERT_SOURCE_RE.match(source_weight_name) + if match is not None: + layer_idx, global_expert_id, projection = match.groups() + global_expert_id = int(global_expert_id) + experts = self.model.layers[int(layer_idx)].mlp.experts + if not experts.fp8_enabled: + raise ValueError( + "Qwen3.6 BF16 checkpoints must use packed 3D expert weights, " + f"got {source_weight_name!r}." + ) + if not experts.is_local_expert(global_expert_id): + return None + return ( + f"model.layers.{layer_idx}.mlp.experts.{global_expert_id}." + f"{projection}.expert_weight" + ) + match = _PACKED_EXPERT_SOURCE_RE.match(source_weight_name) + if match is not None: + layer_idx, projection = match.groups() + experts = self.model.layers[int(layer_idx)].mlp.experts + if experts.fp8_enabled: + raise ValueError( + "Qwen3.6 FP8 checkpoints must use per-expert projections and scales, " + f"got packed tensor {source_weight_name!r}." + ) + return ( + f"model.layers.{layer_idx}.mlp.experts.{projection}." + "packed_expert_weight" + ) + return super().map_weight_name(source_weight_name) + + def resolve_special_weight( + self, + target_weight_name: str, + ) -> WeightTarget | None: + match = _FP8_EXPERT_TARGET_RE.match(target_weight_name) + if match is not None: + layer_idx, global_expert_id, projection = match.groups() + return WeightTarget( + self.model.layers[int(layer_idx)].mlp.experts, + (int(global_expert_id), projection), + ) + match = _PACKED_EXPERT_TARGET_RE.match(target_weight_name) + if match is not None: + layer_idx, projection = match.groups() + return WeightTarget( + self.model.layers[int(layer_idx)].mlp.experts, + projection, + ) + return super().resolve_special_weight(target_weight_name) + + def load_special_weight( + self, + target_weight_name: str, + loaded_weight: torch.Tensor, + loaded_scale: torch.Tensor | None, + ) -> int: + match = _FP8_EXPERT_TARGET_RE.match(target_weight_name) + if match is not None: + layer_idx, global_expert_id, projection = match.groups() + self.model.layers[int(layer_idx)].mlp.experts.load_expert_weight( + int(global_expert_id), + projection, + loaded_weight, + loaded_scale, + ) + return 1 + match = _PACKED_EXPERT_TARGET_RE.match(target_weight_name) + if match is not None: + if loaded_scale is not None: + raise ValueError( + "Qwen3.6 BF16 packed experts unexpectedly have weight scales." + ) + layer_idx, projection = match.groups() + self.model.layers[int(layer_idx)].mlp.experts.load_packed_expert_weight( + projection, + loaded_weight, + ) + return 1 + loaded = super().load_special_weight( + target_weight_name, + loaded_weight, + loaded_scale, + ) + if loaded: + self._loaded_linear_special_weights.add(target_weight_name) + return loaded + + def record_skipped_weight( + self, + source_weight_name: str, + loaded_weight_shape: tuple[int, ...] | None, + loaded_weight_dtype: str | None, + loaded_scale_shape: tuple[int, ...] | None, + loaded_scale_dtype: str | None, + ) -> None: + match = _FP8_EXPERT_SOURCE_RE.match(source_weight_name) + if match is not None: + layer_idx, global_expert_id, projection = match.groups() + experts = self.model.layers[int(layer_idx)].mlp.experts + if not experts.fp8_enabled: + raise ValueError( + f"Qwen3.6 BF16 loader unexpectedly skipped {source_weight_name!r}." + ) + if experts.is_local_expert(int(global_expert_id)): + raise ValueError( + f"Qwen3.6 FP8 loader skipped local expert {source_weight_name!r}." + ) + expected_weight_shape = ( + (experts.hidden_size, experts.global_intermediate_size) + if projection == "down_proj" + else (experts.global_intermediate_size, experts.hidden_size) + ) + expected_scale_shape = ( + (experts.hidden_size // 128, experts.global_intermediate_size // 128) + if projection == "down_proj" + else ( + experts.global_intermediate_size // 128, + experts.hidden_size // 128, + ) + ) + if loaded_weight_shape != expected_weight_shape: + raise ValueError( + "Remote Qwen3.6 FP8 expert weight shape mismatch for " + f"{source_weight_name!r}: expected={expected_weight_shape}, " + f"got={loaded_weight_shape}." + ) + if loaded_weight_dtype != "F8_E4M3": + raise TypeError( + "Remote Qwen3.6 expert weight must be FP8 E4M3, got " + f"safetensors dtype {loaded_weight_dtype}." + ) + if loaded_scale_shape != expected_scale_shape: + raise ValueError( + "Remote Qwen3.6 FP8 expert scale shape mismatch for " + f"{source_weight_name!r}: expected={expected_scale_shape}, " + f"got={loaded_scale_shape}." + ) + if loaded_scale_dtype != "BF16": + raise TypeError( + "Remote Qwen3.6 expert scale must be BF16, got " + f"safetensors dtype {loaded_scale_dtype}." + ) + self._intentionally_skipped_expert_weights.add(source_weight_name) + self._intentionally_skipped_expert_scales.add( + source_weight_name[: -len(".weight")] + ".weight_scale_inv" + ) + return + if not source_weight_name.startswith(self.ignored_weight_prefixes): + raise ValueError( + f"Qwen3.6 MoE loader unexpectedly skipped {source_weight_name!r}." + ) + is_mtp = source_weight_name.startswith("mtp.") + if not is_mtp and ( + loaded_scale_shape is not None or loaded_scale_dtype is not None + ): + raise ValueError( + "Qwen3.6 MoE visual intentional skips must not consume " + f"quantization scales: {source_weight_name!r}." + ) + if is_mtp and loaded_scale_shape is not None: + if loaded_weight_dtype != "F8_E4M3" or loaded_scale_dtype != "BF16": + raise TypeError( + "Qwen3.6 skipped MTP quantized weights require FP8 E4M3 weights " + f"and BF16 scales: {source_weight_name!r}." + ) + if loaded_weight_shape is None or len(loaded_weight_shape) != 2: + raise ValueError( + "Qwen3.6 skipped MTP FP8 weights must be rank-2, " + f"got {source_weight_name!r} shape={loaded_weight_shape}." + ) + expected_scale_shape = tuple( + (int(dimension) + 127) // 128 + for dimension in loaded_weight_shape + ) + if loaded_scale_shape != expected_scale_shape: + raise ValueError( + "Qwen3.6 skipped MTP FP8 scale shape mismatch for " + f"{source_weight_name!r}: expected={expected_scale_shape}, " + f"got={loaded_scale_shape}." + ) + elif is_mtp and loaded_weight_dtype == "F8_E4M3": + raise ValueError( + f"Qwen3.6 skipped MTP FP8 weight is missing its scale: {source_weight_name!r}." + ) + self._intentionally_skipped_weights.add(source_weight_name) + + def validate_loaded_weights(self, loaded_parameter_names: set[str]) -> None: + packed_parameters = { + name + for name, _ in self.named_parameters() + if name.endswith(".mlp.experts.w13_weight") + or name.endswith(".mlp.experts.w2_weight") + } + linear_special_parameters = { + name + for name, _ in self.named_parameters() + if ".linear_attn.in_proj_" in name + and name.endswith(".weight") + and name.rsplit(".", 2)[-2] in {"in_proj_q", "in_proj_k", "in_proj_v"} + } + expected_dense = { + name for name, _ in self.named_parameters() + } - packed_parameters - linear_special_parameters + missing_dense = sorted(expected_dense - loaded_parameter_names) + if missing_dense: + raise ValueError( + f"Missing Qwen3.6 MoE replicated/sharded weights: {missing_dense[:8]}." + ) + + expected_linear_special = { + f"model.layers.{layer_idx}.linear_attn.in_proj_qkv.weight" + for layer_idx in self.config.runtime_layout.linear_attention_layer_indices + } + missing_linear = sorted( + expected_linear_special - self._loaded_linear_special_weights + ) + if missing_linear: + raise ValueError( + f"Missing Qwen3.6 packed GDN weights: {missing_linear[:8]}." + ) + for layer in self.model.layers: + layer.mlp.experts.validate_loaded_weights() + + first_experts = self.model.layers[0].mlp.experts + if first_experts.fp8_enabled: + expected_skipped_experts = { + "model.language_model.layers." + f"{layer_idx}.mlp.experts.{expert_id}.{projection}.weight" + for layer_idx in range(int(self.config.num_hidden_layers)) + for expert_id in range(int(self.config.num_experts)) + if not self.model.layers[layer_idx].mlp.experts.is_local_expert( + expert_id + ) + for projection in ("gate_proj", "up_proj", "down_proj") + } + expected_skipped_scales = { + name[: -len(".weight")] + ".weight_scale_inv" + for name in expected_skipped_experts + } + if self._intentionally_skipped_expert_weights != expected_skipped_experts: + missing = sorted( + expected_skipped_experts + - self._intentionally_skipped_expert_weights + ) + unexpected = sorted( + self._intentionally_skipped_expert_weights + - expected_skipped_experts + ) + raise ValueError( + "Qwen3.6 FP8 remote expert skip audit failed: " + f"missing={missing[:4]}, unexpected={unexpected[:4]}." + ) + if self._intentionally_skipped_expert_scales != expected_skipped_scales: + missing = sorted( + expected_skipped_scales + - self._intentionally_skipped_expert_scales + ) + unexpected = sorted( + self._intentionally_skipped_expert_scales + - expected_skipped_scales + ) + raise ValueError( + "Qwen3.6 FP8 remote expert scale skip audit failed: " + f"missing={missing[:4]}, unexpected={unexpected[:4]}." + ) + + skip_groups = { + "visual": any( + name.startswith(("model.visual.", "visual.")) + for name in self._intentionally_skipped_weights + ), + "mtp": any( + name.startswith("mtp.") + for name in self._intentionally_skipped_weights + ), + } + missing_skip_groups = [name for name, seen in skip_groups.items() if not seen] + if missing_skip_groups: + raise ValueError( + "Qwen3.6 MoE checkpoint is missing expected intentional-skip " + f"groups: {missing_skip_groups}." + ) + logger.info( + "Loaded Qwen3.6 MoE rank {} quantization={} expert_provider={} " + "router_provider={} " + "attention TP {}/{} " + "MoE TP {}/{} EP {}/{} local experts [{}, {}) across {} layers; " + "intentionally skipped {} visual/MTP tensors.", + self.parallel_context.world_rank, + "fp8" if first_experts.fp8_enabled else "bf16", + first_experts.provider.name, + self.model.layers[0].mlp.gate.provider.name, + self.parallel_context.attention_tp_rank, + self.parallel_context.attention_tp_size, + self.parallel_context.moe_tp_rank, + self.parallel_context.moe_tp_size, + self.parallel_context.ep_rank, + self.parallel_context.ep_size, + first_experts.local_expert_start, + first_experts.local_expert_end, + len(self.model.layers), + len(self._intentionally_skipped_weights), + ) diff --git a/src/sparsevllm/models/spec.py b/src/sparsevllm/models/spec.py new file mode 100644 index 00000000..c7547acb --- /dev/null +++ b/src/sparsevllm/models/spec.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sparsevllm.distributed.sharding import validate_model_sharding, validate_top_k +from sparsevllm.distributed.topology import ParallelMode, ParallelTopology +from sparsevllm.utils.config import config_get + + +@dataclass(frozen=True) +class ModelSpec: + name: str + requires_fp8: bool = False + mixed_attention: bool = False + allow_raw_config: bool = False + supports_tiny_random: bool = True + supports_expert_parallel: bool = False + supports_outer_tp_moe: bool = False + supports_data_parallel: bool = False + parallel_mode: ParallelMode = ParallelMode.STANDARD + tiny_random_unquantized: bool = False + skip_checkpoint_validation_in_tiny_random: bool = False + prefix_cache_block_size_multiple: int | None = None + deltakv_checkpoint_model_types: frozenset[str] = frozenset() + runtime_class_name: str = "" + cache_manager_class_name: str = "" + attention_tp_fields: tuple[str, ...] = () + num_experts_field: str | None = None + moe_tp_fields: tuple[str, ...] = () + top_k_field: str | None = None + + def topology(self, tp_size: int, ep_size: int, dp_size: int) -> ParallelTopology: + mode = ( + ParallelMode.OUTER_TP_MOE + if self.supports_outer_tp_moe and int(tp_size) > 1 + else self.parallel_mode + ) + topology = ParallelTopology( + int(tp_size), + int(ep_size), + int(dp_size), + mode, + ) + if topology.expert_parallel_size > 1 and not self.supports_expert_parallel: + raise ValueError( + f"{self.name} does not support expert parallelism, " + f"got EP={topology.expert_parallel_size}." + ) + if topology.data_parallel_size > 1 and not self.supports_data_parallel: + raise ValueError( + f"{self.name} does not support data parallelism, " + f"got DP={topology.data_parallel_size}." + ) + return topology + + def validate_sharding(self, hf_config: Any, topology: ParallelTopology) -> None: + raw_num_experts = ( + config_get(hf_config, self.num_experts_field, None) + if self.num_experts_field + else None + ) + num_experts = int(raw_num_experts) if raw_num_experts is not None else None + validate_model_sharding( + topology, + model_name=self.name, + attention_fields={ + field: int(value) + for field in self.attention_tp_fields + if (value := config_get(hf_config, field, None)) is not None + }, + num_experts=num_experts, + moe_fields={ + field: int(value) + for field in self.moe_tp_fields + if (value := config_get(hf_config, field, None)) is not None + }, + ) + top_k = config_get(hf_config, self.top_k_field, None) if self.top_k_field else None + if top_k is not None and num_experts is not None: + validate_top_k( + self.name, + int(top_k), + num_experts, + ) + + +_DENSE_TP_FIELDS = ( + "num_attention_heads", + "num_key_value_heads", + "vocab_size", + "intermediate_size", +) +_MOE_TP_FIELDS = ("num_attention_heads", "num_key_value_heads", "vocab_size") +_QWEN35_TP_FIELDS = ( + "num_attention_heads", + "num_key_value_heads", + "linear_num_key_heads", + "linear_num_value_heads", + "vocab_size", +) + + +MODEL_SPECS = { + model_type: ModelSpec( + name, + runtime_class_name=runtime_class_name, + attention_tp_fields=_DENSE_TP_FIELDS, + ) + for model_type, (name, runtime_class_name) in { + "qwen2": ("Qwen2", "Qwen2ForCausalLM"), + "qwen3": ("Qwen3", "Qwen3ForCausalLM"), + "llama": ("Llama", "LlamaForCausalLM"), + }.items() +} +MODEL_SPECS.update( + { + "qwen3_5": ModelSpec( + "Qwen3.5", + mixed_attention=True, + allow_raw_config=True, + supports_tiny_random=False, + prefix_cache_block_size_multiple=4096, + deltakv_checkpoint_model_types=frozenset({"qwen3_5"}), + runtime_class_name="Qwen35ForCausalLM", + attention_tp_fields=_QWEN35_TP_FIELDS, + ), + "qwen3_moe": ModelSpec( + "Qwen3MoE", + supports_expert_parallel=True, + supports_outer_tp_moe=True, + runtime_class_name="Qwen3MoeForCausalLM", + attention_tp_fields=_MOE_TP_FIELDS, + num_experts_field="num_experts", + moe_tp_fields=("moe_intermediate_size",), + top_k_field="num_experts_per_tok", + ), + "qwen3_5_moe": ModelSpec( + "Qwen3.6 MoE", + mixed_attention=True, + allow_raw_config=True, + supports_tiny_random=False, + supports_expert_parallel=True, + supports_outer_tp_moe=True, + prefix_cache_block_size_multiple=4096, + deltakv_checkpoint_model_types=frozenset({"qwen3_5"}), + runtime_class_name="Qwen35MoeForCausalLM", + attention_tp_fields=( + *_QWEN35_TP_FIELDS, + "shared_expert_intermediate_size", + ), + num_experts_field="num_experts", + moe_tp_fields=("moe_intermediate_size",), + top_k_field="num_experts_per_tok", + ), + "minimax_m2": ModelSpec( + "MiniMax M2.7", + requires_fp8=True, + supports_expert_parallel=True, + supports_outer_tp_moe=True, + runtime_class_name="MiniMaxM2ForCausalLM", + attention_tp_fields=_MOE_TP_FIELDS, + num_experts_field="num_local_experts", + moe_tp_fields=("intermediate_size",), + top_k_field="num_experts_per_tok", + ), + "deepseek_v4": ModelSpec( + "DeepSeek V4 Flash", + requires_fp8=True, + supports_expert_parallel=True, + supports_data_parallel=True, + parallel_mode=ParallelMode.DPA_EP, + tiny_random_unquantized=True, + skip_checkpoint_validation_in_tiny_random=True, + runtime_class_name="DeepseekV4ForCausalLM", + cache_manager_class_name="DeepseekV4CacheManager", + attention_tp_fields=("num_attention_heads", "num_key_value_heads"), + num_experts_field="n_routed_experts", + moe_tp_fields=("moe_intermediate_size",), + top_k_field="num_experts_per_tok", + ), + } +) + + +def resolve_model_spec(model_type: str) -> ModelSpec: + if model_type not in MODEL_SPECS: + supported = ", ".join(sorted(MODEL_SPECS)) + raise NotImplementedError( + f"Unsupported Sparse-vLLM model_type={model_type!r}; supported: {supported}." + ) + return MODEL_SPECS[model_type] diff --git a/src/sparsevllm/operators/all_reduce.py b/src/sparsevllm/operators/all_reduce.py new file mode 100644 index 00000000..ee848e3d --- /dev/null +++ b/src/sparsevllm/operators/all_reduce.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from typing import Protocol + +import torch +import torch.distributed as dist + +from sparsevllm.utils.log import logger + + +class AllReduceProvider(Protocol): + name: str + + def run(self, tensor: torch.Tensor) -> torch.Tensor: ... + + +class TorchDistributedAllReduceProvider: + name = "torch_distributed" + + def __init__(self, group: dist.ProcessGroup | None) -> None: + self.group = group + + def run(self, tensor: torch.Tensor) -> torch.Tensor: + dist.all_reduce(tensor, group=self.group) + return tensor + + +class HopperTp2FlashInferAllReduceProvider: + name = "hopper_tp2_flashinfer" + hidden_size = 2048 + max_rows = 256 + + def __init__(self, group: dist.ProcessGroup) -> None: + from flashinfer import comm + + self.comm = comm + self.fallback = TorchDistributedAllReduceProvider(group) + self.workspace = comm.create_allreduce_fusion_workspace( + backend="trtllm", + world_size=2, + rank=dist.get_rank(group), + max_token_num=self.max_rows, + hidden_dim=self.hidden_size, + dtype=torch.bfloat16, + group=group, + ) + + def _supports(self, tensor: torch.Tensor) -> bool: + return ( + tensor.is_cuda + and tensor.dtype == torch.bfloat16 + and tensor.is_contiguous() + and tensor.ndim >= 2 + and tensor.shape[-1] == self.hidden_size + and tensor.numel() <= self.max_rows * self.hidden_size + ) + + def run(self, tensor: torch.Tensor) -> torch.Tensor: + if not self._supports(tensor): + return self.fallback.run(tensor) + output = self.comm.allreduce_fusion( + input=tensor.view(-1, self.hidden_size), + workspace=self.workspace, + pattern=self.comm.AllReduceFusionPattern.kAllReduce, + launch_with_pdl=True, + trigger_completion_at_end=tensor.numel() > 16 * self.hidden_size, + ) + return output.view_as(tensor) + + +def resolve_all_reduce_provider( + group: dist.ProcessGroup | None, + world_size: int, +) -> AllReduceProvider: + if ( + world_size == 2 + and group is not None + and dist.get_backend(group) == dist.Backend.NCCL + and torch.cuda.get_device_capability() == (9, 0) + ): + provider = HopperTp2FlashInferAllReduceProvider(group) + logger.info("AllReduce provider: %s", provider.name) + return provider + return TorchDistributedAllReduceProvider(group) diff --git a/src/sparsevllm/operators/fp8_linear.py b/src/sparsevllm/operators/fp8_linear.py index 068d9ca6..0ba23627 100644 --- a/src/sparsevllm/operators/fp8_linear.py +++ b/src/sparsevllm/operators/fp8_linear.py @@ -2,7 +2,6 @@ from dataclasses import dataclass from importlib.util import find_spec -from shutil import which import torch @@ -14,7 +13,6 @@ runtime_version_at_least, ) from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum -from sparsevllm.utils.log import logger @dataclass(frozen=True) @@ -53,9 +51,6 @@ class FlashInferSm90Fp8LinearProvider(Fp8LinearProvider): name = "flashinfer_sm90" priority = 100 - def __init__(self) -> None: - self._fallback: TritonFp8LinearProvider | None = None - @classmethod def supports(cls, spec: Fp8LinearSpec, caps: DeviceCaps) -> SupportResult: if spec.block_shape != (128, 128): @@ -88,8 +83,6 @@ def supports(cls, spec: Fp8LinearSpec, caps: DeviceCaps) -> SupportResult: return SupportResult.yes() def __call__(self, x, weight, weight_scale_inv, bias=None): - if self._fallback is not None: - return self._fallback(x, weight, weight_scale_inv, bias) if x.dtype != torch.bfloat16: raise TypeError( f"FlashInfer SM90 FP8 Linear requires BF16 activations, got {x.dtype}." @@ -97,32 +90,12 @@ def __call__(self, x, weight, weight_scale_inv, bias=None): from flashinfer.gemm import fp8_blockscale_gemm_sm90 original_shape = x.shape[:-1] - try: - output = fp8_blockscale_gemm_sm90( - x.reshape(-1, x.shape[-1]).contiguous(), - weight, - weight_scale=weight_scale_inv, - out_dtype=torch.bfloat16, - ) - except RuntimeError as exc: - message = str(exc) - missing_jit_artifact = ( - which("nvcc") is None - and ( - "!cubin.empty() || isPathValid(path_)" in message - or "nvcc not found" in message.lower() - or "could not find nvcc" in message.lower() - ) - ) - if not missing_jit_artifact: - raise - self._fallback = TritonFp8LinearProvider() - logger.warning( - "FlashInfer SM90 FP8 Linear has no cached kernel for " - f"shape={tuple(x.shape)}x{tuple(weight.shape)} and nvcc is " - "unavailable; binding this Linear instance to Triton." - ) - return self._fallback(x, weight, weight_scale_inv, bias) + output = fp8_blockscale_gemm_sm90( + x.reshape(-1, x.shape[-1]).contiguous(), + weight, + weight_scale=weight_scale_inv, + out_dtype=torch.bfloat16, + ) if bias is not None: output.add_(bias) return output.reshape(*original_shape, weight.shape[0]) diff --git a/src/sparsevllm/operators/gate_up_swiglu.py b/src/sparsevllm/operators/gate_up_swiglu.py new file mode 100644 index 00000000..c4ac8c6a --- /dev/null +++ b/src/sparsevllm/operators/gate_up_swiglu.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +from dataclasses import dataclass +import torch +import torch.nn.functional as F + +import sparsevllm.platforms as platforms +from sparsevllm.operators.registry import OpRegistry, OpResolver, SupportResult +from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum + + +@dataclass(frozen=True) +class GateUpSwiGLUOpSpec: + hidden_size: int + intermediate_size: int + tp_size: int + activation_dtype: torch.dtype + weight_dtype: torch.dtype + cuda_graph: bool + + def __post_init__(self) -> None: + if min(self.hidden_size, self.intermediate_size, self.tp_size) <= 0: + raise ValueError("Gate/up SwiGLU dimensions and TP size must be positive.") + if self.intermediate_size % self.tp_size: + raise ValueError( + "Gate/up SwiGLU intermediate size must be divisible by TP size." + ) + if not ( + self.activation_dtype.is_floating_point + and self.weight_dtype.is_floating_point + ): + raise TypeError( + "Gate/up SwiGLU activations and weights must be floating point." + ) + + +class GateUpSwiGLUProvider: + name = "" + priority = 0 + + def run( + self, + spec: GateUpSwiGLUOpSpec, + inputs: torch.Tensor, + projection, + ) -> torch.Tensor: + raise NotImplementedError + + +GATE_UP_SWIGLU_REGISTRY: OpRegistry[ + GateUpSwiGLUOpSpec, GateUpSwiGLUProvider +] = OpRegistry("gate/up SwiGLU") + + +@GATE_UP_SWIGLU_REGISTRY.register +class NativeGateUpSwiGLUProvider(GateUpSwiGLUProvider): + name = "native" + priority = 0 + + @classmethod + def supports( + cls, + spec: GateUpSwiGLUOpSpec, + caps: DeviceCaps, + ) -> SupportResult: + del spec, caps + return SupportResult.yes() + + def run( + self, + spec: GateUpSwiGLUOpSpec, + inputs: torch.Tensor, + projection, + ) -> torch.Tensor: + del spec + gate, up = projection(inputs).chunk(2, dim=-1) + return F.silu(gate, inplace=True).mul_(up) + + +@GATE_UP_SWIGLU_REGISTRY.register +class H20GateUpSwiGLUProvider(NativeGateUpSwiGLUProvider): + name = "h20_triton_decode" + priority = 20 + + @classmethod + def supports( + cls, + spec: GateUpSwiGLUOpSpec, + caps: DeviceCaps, + ) -> SupportResult: + if caps.platform != PlatformEnum.CUDA or caps.compute_capability != (9, 0): + return SupportResult.no( + f"requires CUDA SM90, got {caps.platform.name} " + f"{caps.compute_capability}" + ) + if caps.device_name != "NVIDIA H20": + return SupportResult.no( + f"requires profiled NVIDIA H20 hardware, got {caps.device_name}" + ) + if not caps.supports_triton: + return SupportResult.no("platform does not support Triton") + if not caps.supports_bfloat16: + return SupportResult.no("device does not support BF16") + if spec.cuda_graph and not caps.supports_graph_capture: + return SupportResult.no("device does not support CUDA Graph capture") + if spec.activation_dtype != torch.bfloat16 or spec.weight_dtype != torch.bfloat16: + return SupportResult.no( + "requires BF16 activations and weights, got " + f"{spec.activation_dtype} and {spec.weight_dtype}" + ) + if (spec.hidden_size, spec.intermediate_size, spec.tp_size) not in { + (2048, 512, 1), + (2048, 512, 2), + }: + return SupportResult.no( + "requires profiled (hidden, intermediate, TP) shape in " + "{(2048, 512, 1), (2048, 512, 2)}" + ) + return SupportResult.yes() + + def run( + self, + spec: GateUpSwiGLUOpSpec, + inputs: torch.Tensor, + projection, + ) -> torch.Tensor: + if inputs.shape[0] != 1: + return super().run(spec, inputs, projection) + from sparsevllm.triton_kernel.gate_up_swiglu import h20_gate_up_swiglu + + return h20_gate_up_swiglu(inputs, projection.weight) + + +def resolve_gate_up_swiglu_provider( + spec: GateUpSwiGLUOpSpec, + *, + device_index: int | None = None, +) -> GateUpSwiGLUProvider: + platform = platforms.current_platform + if device_index is None: + device_index = torch.cuda.current_device() if platform.is_cuda_alike() else 0 + caps = platform.get_device_caps(int(device_index)) + return OpResolver(GATE_UP_SWIGLU_REGISTRY).resolve(spec, caps).provider diff --git a/src/sparsevllm/operators/gated_shared_add.py b/src/sparsevllm/operators/gated_shared_add.py new file mode 100644 index 00000000..df9de374 --- /dev/null +++ b/src/sparsevllm/operators/gated_shared_add.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import torch + + +def gated_shared_add( + routed: torch.Tensor, + shared: torch.Tensor, + gate_logits: torch.Tensor, +) -> torch.Tensor: + if ( + routed.ndim != 2 + or routed.shape != shared.shape + or gate_logits.shape != (routed.shape[0], 1) + ): + raise ValueError( + "gated_shared_add expects routed/shared [tokens, hidden] and gate " + f"[tokens, 1], got {tuple(routed.shape)}, {tuple(shared.shape)}, " + f"{tuple(gate_logits.shape)}." + ) + if ( + routed.dtype != torch.bfloat16 + or shared.dtype != routed.dtype + or gate_logits.dtype != routed.dtype + ): + raise TypeError("gated_shared_add requires BF16 inputs with matching dtypes.") + if ( + not routed.is_cuda + or shared.device != routed.device + or gate_logits.device != routed.device + ): + raise ValueError("gated_shared_add requires CUDA inputs on one device.") + if ( + not routed.is_contiguous() + or not shared.is_contiguous() + or gate_logits.stride(1) != 1 + ): + raise ValueError("gated_shared_add requires contiguous hidden dimensions.") + + from sparsevllm.triton_kernel.qwen3_5.gated_shared_add import ( + triton_gated_shared_add, + ) + + return triton_gated_shared_add(routed, shared, gate_logits) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index e0d212fd..0a1d5dc3 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -30,6 +30,7 @@ class MoeOpSpec: tp_size: int = 1 routing_method: str = "softmax" scale_dtype: torch.dtype | None = None + activation_limit: float | None = None def __post_init__(self) -> None: if self.num_experts <= 0 or self.num_local_experts <= 0: @@ -351,11 +352,173 @@ def run( return output +@MOE_REGISTRY.register +class FlashInferCutlassFp4MoeProvider(MoeProvider): + """Hopper W4A16 MoE for DeepSeek V4's packed MXFP4 experts.""" + + name = "flashinfer_cutlass_fp4_sm90" + priority = 105 + gate_up_order = "up_gate" + + @classmethod + def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: + if spec.tp_size != 1: + return SupportResult.no("does not support tensor-parallel expert shards") + if spec.weight_dtype != torch.uint8: + return SupportResult.no(f"requires packed FP4 uint8 weights, got {spec.weight_dtype}") + if spec.scale_dtype != torch.uint8: + return SupportResult.no(f"requires raw UE8M0 uint8 scales, got {spec.scale_dtype}") + if spec.block_shape != (1, 32): + return SupportResult.no(f"requires per-row K32 scales, got {spec.block_shape}") + if spec.hidden_size % 128 or spec.intermediate_size % 128: + return SupportResult.no("FP4 hidden/intermediate sizes must be 128-aligned") + if spec.activation_dtype != torch.bfloat16: + return SupportResult.no( + f"requires BF16 activations, got {spec.activation_dtype}" + ) + if caps.platform != PlatformEnum.CUDA or caps.compute_capability != (9, 0): + return SupportResult.no( + f"requires CUDA SM90, got {caps.platform.name} {caps.compute_capability}" + ) + if not runtime_version_at_least(caps.runtime_version, (12, 8)): + return SupportResult.no( + "requires CUDA runtime >= 12.8, " + f"got {caps.runtime_version or 'unknown'}" + ) + if find_spec("flashinfer") is None: + return SupportResult.no("flashinfer is not installed") + return SupportResult.yes() + + def load_expert_projection( + self, + spec, + *, + local_expert_id, + projection, + loaded_weight, + loaded_scale, + w13_weight, + w2_weight, + w13_scale_inv, + w2_scale_inv, + ) -> None: + if loaded_scale is None or w13_scale_inv is None or w2_scale_inv is None: + raise RuntimeError("DeepSeek V4 FP4 experts require UE8M0 scales.") + if projection == "down": + weight_target = w2_weight[local_expert_id] + scale_target = w2_scale_inv[local_expert_id] + elif projection in {"gate", "up"}: + offset = self._packed_projection_offset(projection, spec.intermediate_size) + weight_target = w13_weight[ + local_expert_id, offset : offset + spec.intermediate_size + ] + scale_target = w13_scale_inv[ + local_expert_id, offset : offset + spec.intermediate_size + ] + else: + raise ValueError(f"Unknown logical MoE projection {projection!r}.") + if tuple(loaded_weight.shape) != tuple(weight_target.shape): + raise ValueError( + "Packed FP4 expert weight shape mismatch: " + f"expected={tuple(weight_target.shape)}, got={tuple(loaded_weight.shape)}." + ) + if tuple(loaded_scale.shape) != tuple(scale_target.shape): + raise ValueError( + "FP4 expert scale shape mismatch: " + f"expected={tuple(scale_target.shape)}, got={tuple(loaded_scale.shape)}." + ) + weight_target.copy_(loaded_weight.contiguous().view(torch.uint8)) + scale_target.copy_(loaded_scale.contiguous().view(torch.uint8)) + + @staticmethod + def prepare_weights( + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + from flashinfer import fused_moe + + return ( + fused_moe.interleave_moe_weights_for_sm90_mixed_gemm( + w13_weight.contiguous(), "fp4" + ), + fused_moe.interleave_moe_weights_for_sm90_mixed_gemm( + w2_weight.contiguous(), "fp4" + ), + fused_moe.interleave_moe_scales_for_sm90_mixed_gemm( + w13_scale.contiguous() + ).view(torch.int32), + fused_moe.interleave_moe_scales_for_sm90_mixed_gemm( + w2_scale.contiguous() + ).view(torch.int32), + ) + + def run( + self, + spec, + hidden_states, + topk_ids, + topk_weights, + w13_weight, + w2_weight, + w13_scale_inv, + w2_scale_inv, + *, + local_expert_start, + ep_rank, + ): + del local_expert_start + if w13_scale_inv is None or w2_scale_inv is None: + raise RuntimeError("FlashInfer FP4 MoE requires prepared expert scales.") + if w13_scale_inv.dtype != torch.int32 or w2_scale_inv.dtype != torch.int32: + raise RuntimeError("FP4 expert weights must be prepared before inference.") + from flashinfer.fused_moe import cutlass_fused_moe + from flashinfer.tllm_enums import ActivationType + + output = torch.empty_like(hidden_states) + activation_limit = float(spec.activation_limit or 0.0) + if activation_limit <= 0: + raise RuntimeError("DeepSeek V4 FP4 MoE requires a positive activation limit.") + alpha = torch.ones( + spec.num_local_experts, dtype=torch.float32, device=hidden_states.device + ) + beta = torch.zeros_like(alpha) + limit = torch.full_like(alpha, activation_limit) + cutlass_fused_moe( + hidden_states, + topk_ids.to(dtype=torch.int32), + topk_weights.to(dtype=torch.float32), + w13_weight, + w2_weight, + hidden_states.dtype, + quant_scales=[w13_scale_inv, w2_scale_inv], + swiglu_alpha=alpha, + swiglu_beta=beta, + swiglu_limit=limit, + ep_size=int(spec.ep_size), + ep_rank=int(ep_rank), + output=output, + use_w4_group_scaling=True, + use_fused_finalize=False, + enable_pdl=False, + activation_type=ActivationType.Swiglu, + ) + return output + + @MOE_REGISTRY.register class TritonHopperFusedMoeProvider(MoeProvider): name = "triton_hopper_fused" priority = 20 gate_up_order = "gate_up" + PROFILED_DEVICE_NAME = "NVIDIA H100 80GB HBM3" + PROFILED_SHAPES = ( + (128, 64, 2048, 384, 8, 2, 2), + (256, 256, 2048, 512, 8, 1, 1), + (256, 256, 2048, 256, 8, 2, 1), + (256, 128, 2048, 512, 8, 1, 2), + ) @classmethod def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: @@ -373,14 +536,13 @@ def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: ) if spec.weight_dtype != torch.bfloat16 or spec.block_shape is not None: return SupportResult.no("requires unquantized BF16 expert weights") - if caps.device_name != "NVIDIA H100 80GB HBM3": + if caps.device_name != cls.PROFILED_DEVICE_NAME: return SupportResult.no( - "requires profiled NVIDIA H100 80GB HBM3 hardware, " + f"requires profiled {cls.PROFILED_DEVICE_NAME} hardware, " f"got {caps.device_name}" ) if not caps.supports_bfloat16: return SupportResult.no("device does not support BF16") - profiled_shape = (128, 64, 2048, 384, 8, 2, 2) actual_shape = ( spec.num_experts, spec.num_local_experts, @@ -390,10 +552,10 @@ def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: spec.tp_size, spec.ep_size, ) - if actual_shape != profiled_shape: + if actual_shape not in cls.PROFILED_SHAPES: return SupportResult.no( - "requires profiled TP2xEP2 MoE shape " - f"{profiled_shape}, got {actual_shape}" + "requires a profiled MoE shape in " + f"{cls.PROFILED_SHAPES}, got {actual_shape}" ) return SupportResult.yes() @@ -427,6 +589,18 @@ def run( ) +@MOE_REGISTRY.register +class H20Qwen36FusedMoeProvider(TritonHopperFusedMoeProvider): + name = "h20_qwen36_fused_bf16" + priority = 21 + PROFILED_DEVICE_NAME = "NVIDIA H20" + PROFILED_SHAPES = ( + (256, 256, 2048, 512, 8, 1, 1), + (256, 256, 2048, 256, 8, 2, 1), + (256, 128, 2048, 512, 8, 1, 2), + ) + + @MOE_REGISTRY.register class TritonMoeProvider(MoeProvider): name = "triton" @@ -504,6 +678,112 @@ def run( ) +@MOE_REGISTRY.register +class HopperQwen36HybridFp8MoeProvider(FlashInferCutlassFp8MoeProvider): + """Bind one weight layout and dispatch profiled token buckets by kernel.""" + + name = "hopper_qwen36_hybrid_fp8" + priority = 110 + PROFILED_DEVICE_NAME = "NVIDIA H100 80GB HBM3" + PROFILED_SHAPES = frozenset( + { + (256, 256, 2048, 512, 8, 1, 1), + (256, 128, 2048, 512, 8, 1, 2), + } + ) + TRITON_MAX_TOKENS_BY_EP_SIZE = {1: 8, 2: 4} + + @classmethod + def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: + if spec.cuda_graph and not caps.supports_graph_capture: + return SupportResult.no("device does not support CUDA Graph capture") + if caps.device_name != cls.PROFILED_DEVICE_NAME: + return SupportResult.no( + f"requires profiled {cls.PROFILED_DEVICE_NAME} hardware, " + f"got {caps.device_name}" + ) + actual_shape = ( + spec.num_experts, + spec.num_local_experts, + spec.hidden_size, + spec.intermediate_size, + spec.top_k, + spec.tp_size, + spec.ep_size, + ) + if actual_shape not in cls.PROFILED_SHAPES: + return SupportResult.no( + "requires profiled Qwen3.6 TP1/EP1 or " + "global-TP2/MoE-TP1xEP2 shape " + f"{sorted(cls.PROFILED_SHAPES)}, " + f"got {actual_shape}" + ) + flashinfer_support = super().supports(spec, caps) + if not flashinfer_support.supported: + return SupportResult.no( + f"FlashInfer prefill path: {flashinfer_support.reason}" + ) + triton_support = TritonMoeProvider.supports(spec, caps) + if not triton_support.supported: + return SupportResult.no( + f"Triton decode path: {triton_support.reason}" + ) + return SupportResult.yes() + + def run( + self, + spec, + hidden_states, + topk_ids, + topk_weights, + w13_weight, + w2_weight, + w13_scale_inv, + w2_scale_inv, + *, + local_expert_start, + ep_rank, + ): + triton_max_tokens = self.TRITON_MAX_TOKENS_BY_EP_SIZE[int(spec.ep_size)] + if int(hidden_states.shape[0]) > triton_max_tokens: + return super().run( + spec, + hidden_states, + topk_ids, + topk_weights, + w13_weight, + w2_weight, + w13_scale_inv, + w2_scale_inv, + local_expert_start=local_expert_start, + ep_rank=ep_rank, + ) + if w13_scale_inv is None or w2_scale_inv is None: + raise RuntimeError("Qwen3.6 hybrid FP8 MoE requires expert scales.") + from sparsevllm.triton_kernel.moe import fused_moe_fp8 + + return fused_moe_fp8( + hidden_states, + w13_weight, + w2_weight, + w13_scale_inv, + w2_scale_inv, + topk_ids, + topk_weights, + num_experts=spec.num_experts, + local_expert_start=local_expert_start, + gate_up_order=self.gate_up_order, + ) + + +@MOE_REGISTRY.register +class H20Qwen36HybridFp8MoeProvider(HopperQwen36HybridFp8MoeProvider): + name = "h20_qwen36_hybrid_fp8" + priority = 111 + PROFILED_DEVICE_NAME = "NVIDIA H20" + TRITON_MAX_TOKENS_BY_EP_SIZE = {1: 8, 2: 1} + + def resolve_moe_provider( spec: MoeOpSpec, *, diff --git a/src/sparsevllm/operators/moe_router.py b/src/sparsevllm/operators/moe_router.py new file mode 100644 index 00000000..fc8e74c3 --- /dev/null +++ b/src/sparsevllm/operators/moe_router.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +import sparsevllm.platforms as platforms +from sparsevllm.operators.registry import ( + OpRegistry, + OpResolver, + SupportResult, +) +from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum + + +@dataclass(frozen=True) +class MoeRouterOpSpec: + num_experts: int + top_k: int + activation_dtype: torch.dtype + norm_topk_prob: bool + cuda_graph: bool + + def __post_init__(self) -> None: + if self.num_experts <= 0: + raise ValueError("MoE router num_experts must be positive.") + if not 1 <= self.top_k <= self.num_experts: + raise ValueError( + f"MoE router top_k must be in [1, {self.num_experts}], " + f"got {self.top_k}." + ) + if not self.activation_dtype.is_floating_point: + raise TypeError( + "MoE router activations must be floating point, " + f"got {self.activation_dtype}." + ) + + +class MoeRouterProvider: + name = "" + priority = 0 + + def run( + self, + spec: MoeRouterOpSpec, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError + + +MOE_ROUTER_REGISTRY: OpRegistry[MoeRouterOpSpec, MoeRouterProvider] = OpRegistry( + "MoE router" +) + + +@MOE_ROUTER_REGISTRY.register +class TritonMoeRouterProvider(MoeRouterProvider): + name = "triton" + priority = 10 + + @classmethod + def supports( + cls, + spec: MoeRouterOpSpec, + caps: DeviceCaps, + ) -> SupportResult: + if caps.platform != PlatformEnum.CUDA: + return SupportResult.no(f"requires CUDA, got {caps.platform.name}") + if not caps.supports_triton: + return SupportResult.no("platform does not support Triton") + if spec.cuda_graph and not caps.supports_graph_capture: + return SupportResult.no("device does not support CUDA Graph capture") + if spec.activation_dtype not in {torch.bfloat16, torch.float16}: + return SupportResult.no( + f"requires BF16 or FP16 logits, got {spec.activation_dtype}" + ) + if spec.num_experts not in {128, 256} or spec.top_k != 8: + return SupportResult.no( + "requires num_experts in {128, 256} and top_k=8" + ) + return SupportResult.yes() + + def run( + self, + spec: MoeRouterOpSpec, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + from sparsevllm.triton_kernel.moe_topk import topk_softmax + + return topk_softmax( + router_logits, + top_k=spec.top_k, + norm_topk_prob=spec.norm_topk_prob, + ) + + +def resolve_moe_router_provider( + spec: MoeRouterOpSpec, + *, + device_index: int | None = None, +) -> MoeRouterProvider: + platform = platforms.current_platform + if device_index is None: + device_index = torch.cuda.current_device() if platform.is_cuda_alike() else 0 + caps = platform.get_device_caps(int(device_index)) + return OpResolver(MOE_ROUTER_REGISTRY).resolve(spec, caps).provider diff --git a/src/sparsevllm/operators/registry.py b/src/sparsevllm/operators/registry.py index e174d301..2081ac7c 100644 --- a/src/sparsevllm/operators/registry.py +++ b/src/sparsevllm/operators/registry.py @@ -1,15 +1,43 @@ from __future__ import annotations +import weakref from dataclasses import dataclass from typing import Generic, Protocol, TypeVar from sparsevllm.platforms.interface import DeviceCaps +from sparsevllm.utils.log import logger SpecT = TypeVar("SpecT") ProviderT = TypeVar("ProviderT", bound="OperatorProvider") +_OPERATOR_BINDINGS: dict[str, weakref.WeakSet[object]] = {} + + +def record_operator_binding(operator_type: str, provider: object) -> None: + _OPERATOR_BINDINGS.setdefault(operator_type, weakref.WeakSet()).add(provider) + + +def _implementation_name(provider: object) -> str: + return getattr(provider, "implementation_name", None) or getattr(provider, "name", None) or provider.provider_name + + +def log_operator_implementations() -> None: + entries = sorted( + ( + operator_type, + ", ".join(sorted({_implementation_name(provider) for provider in providers})), + ) + for operator_type, providers in _OPERATOR_BINDINGS.items() + if providers + ) + if not entries: + return + rows = "\n".join(f" {operator_type}: {implementation}" for operator_type, implementation in entries) + logger.info("Operator implementations:\n{}", rows) + + def runtime_version_at_least( version: str | None, minimum: tuple[int, int], @@ -99,4 +127,5 @@ def resolve( ) supported.sort(key=lambda provider: (-int(provider.priority), provider.name)) selected = supported[0](**provider_kwargs) + record_operator_binding(self.registry.family, selected) return ResolvedProvider(selected, tuple(rejected)) diff --git a/src/sparsevllm/quantization/__init__.py b/src/sparsevllm/quantization/__init__.py index 91bb131e..263d204e 100644 --- a/src/sparsevllm/quantization/__init__.py +++ b/src/sparsevllm/quantization/__init__.py @@ -1,7 +1,9 @@ +from sparsevllm.quantization.config import QuantizationConfig from sparsevllm.operators.fp8_linear import resolve_fp8_linear_provider from sparsevllm.quantization.registry import QuantizationRegistry __all__ = [ "QuantizationRegistry", + "QuantizationConfig", "resolve_fp8_linear_provider", ] diff --git a/src/sparsevllm/quantization/config.py b/src/sparsevllm/quantization/config.py new file mode 100644 index 00000000..4a77238d --- /dev/null +++ b/src/sparsevllm/quantization/config.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sparsevllm.utils.config import config_get + + +@dataclass(frozen=True) +class QuantizationConfig: + enabled: bool = False + quant_method: str = "" + weight_dtype: str = "" + activation_scheme: str = "" + weight_block_size: tuple[int, int] | None = None + model_name: str = "qwen3_5" + + @classmethod + def disabled(cls, *, model_name: str = "qwen3_5") -> "QuantizationConfig": + return cls(model_name=model_name) + + def to_dict(self) -> dict[str, Any]: + if not self.enabled: + return {} + payload: dict[str, Any] = { + "quant_method": self.quant_method, + "fmt": self.weight_dtype, + "activation_scheme": self.activation_scheme, + } + if self.weight_block_size is not None: + payload["weight_block_size"] = list(self.weight_block_size) + return payload + + @classmethod + def from_hf_config( + cls, + value: Any, + *, + required_fp8: bool = False, + model_name: str = "qwen3_5", + ) -> "QuantizationConfig": + if value is None: + if required_fp8: + raise ValueError( + f"{model_name} requires FP8 quantization_config; " + "BF16/FP16 fallback is not supported." + ) + return cls.disabled(model_name=model_name) + + quant_method = str( + config_get(value, "quant_method", config_get(value, "method", "")) + or "" + ).strip().lower() + if quant_method not in {"fp8", "fbgemm_fp8"}: + if required_fp8: + raise ValueError( + f"{model_name} requires quantization_config.quant_method='fp8', " + f"got {quant_method!r}." + ) + if quant_method: + raise NotImplementedError( + f"Sparse-vLLM does not support quant_method={quant_method!r} " + f"for {model_name}." + ) + return cls.disabled(model_name=model_name) + + weight_dtype = str( + config_get( + value, + "weight_dtype", + config_get(value, "fmt", config_get(value, "format", "e4m3")), + ) + or "" + ).strip().lower() + if "e4m3" not in weight_dtype: + raise ValueError( + f"Sparse-vLLM {model_name} FP8 supports e4m3 weights only, " + f"got weight_dtype={weight_dtype!r}." + ) + + activation_scheme = str( + config_get( + value, + "activation_scheme", + config_get(value, "activation", "dynamic"), + ) + or "" + ).strip().lower() + if activation_scheme != "dynamic": + raise ValueError( + f"Sparse-vLLM {model_name} FP8 supports dynamic activation only, " + f"got activation_scheme={activation_scheme!r}." + ) + + block_size = config_get( + value, + "weight_block_size", + config_get( + value, + "weight_block_shape", + config_get(value, "block_size", (128, 128)), + ), + ) + if isinstance(block_size, int): + block_tuple = (int(block_size), int(block_size)) + elif isinstance(block_size, (list, tuple)) and len(block_size) == 2: + block_tuple = (int(block_size[0]), int(block_size[1])) + else: + raise ValueError(f"weight_block_size must be a pair, got {block_size!r}.") + if block_tuple != (128, 128): + raise ValueError( + f"Sparse-vLLM {model_name} FP8 supports " + "weight_block_size=(128, 128) only, " + f"got {block_tuple}." + ) + + return cls( + enabled=True, + quant_method="fp8", + weight_dtype="e4m3", + activation_scheme="dynamic", + weight_block_size=block_tuple, + model_name=model_name, + ) diff --git a/src/sparsevllm/triton_kernel/gate_up_swiglu.py b/src/sparsevllm/triton_kernel/gate_up_swiglu.py new file mode 100644 index 00000000..8cb85370 --- /dev/null +++ b/src/sparsevllm/triton_kernel/gate_up_swiglu.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +_H20_DECODE_CONFIGS = { + (2048, 256): dict(block_m=16, block_n=32, block_k=64, warps=4, stages=4), + (2048, 512): dict(block_m=16, block_n=32, block_k=64, warps=4, stages=4), +} + + +@triton.jit +def _gate_up_swiglu_kernel( + input_ptr, + weight_ptr, + output_ptr, + N: tl.constexpr, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + m_offsets = tl.arange(0, BLOCK_M) + n_offsets = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + k_offsets = tl.arange(0, BLOCK_K) + input_ptrs = input_ptr + m_offsets[:, None] * K + k_offsets[None, :] + gate_ptrs = weight_ptr + n_offsets[None, :] * K + k_offsets[:, None] + up_ptrs = gate_ptrs + N * K + gate_accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + up_accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + for k_start in range(0, tl.cdiv(K, BLOCK_K)): + remaining_k = K - k_start * BLOCK_K + input_values = tl.load( + input_ptrs, + mask=(m_offsets[:, None] == 0) + & (k_offsets[None, :] < remaining_k), + other=0.0, + ) + weight_mask = (k_offsets[:, None] < remaining_k) & ( + n_offsets[None, :] < N + ) + gate_accumulator += tl.dot( + input_values, + tl.load(gate_ptrs, mask=weight_mask, other=0.0), + ) + up_accumulator += tl.dot( + input_values, + tl.load(up_ptrs, mask=weight_mask, other=0.0), + ) + input_ptrs += BLOCK_K + gate_ptrs += BLOCK_K + up_ptrs += BLOCK_K + + element_dtype = weight_ptr.dtype.element_ty + gate = gate_accumulator.to(element_dtype).to(tl.float32) + up = up_accumulator.to(element_dtype) + gate = (gate / (1.0 + tl.exp(-gate))).to(element_dtype) + tl.store( + output_ptr + m_offsets[:, None] * N + n_offsets[None, :], + gate * up, + mask=(m_offsets[:, None] == 0) & (n_offsets[None, :] < N), + ) + + +def h20_gate_up_swiglu(inputs: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + if inputs.ndim != 2 or weight.ndim != 2 or weight.shape[0] % 2: + raise ValueError( + "h20_gate_up_swiglu expects input [1, K] and weight [2N, K]." + ) + shape = (int(inputs.shape[1]), int(weight.shape[0]) // 2) + config = _H20_DECODE_CONFIGS.get(shape) + if inputs.shape != (1, shape[0]) or weight.shape != (2 * shape[1], shape[0]) or config is None: + raise ValueError(f"No H20 gate/up SwiGLU config for shape {shape}.") + if inputs.dtype != torch.bfloat16 or weight.dtype != inputs.dtype: + raise TypeError("h20_gate_up_swiglu requires matching BF16 tensors.") + if not inputs.is_cuda or weight.device != inputs.device: + raise ValueError("h20_gate_up_swiglu requires CUDA tensors on one device.") + if not inputs.is_contiguous() or not weight.is_contiguous(): + raise ValueError("h20_gate_up_swiglu requires contiguous tensors.") + + output = torch.empty((1, shape[1]), dtype=inputs.dtype, device=inputs.device) + _gate_up_swiglu_kernel[(triton.cdiv(shape[1], config["block_n"]),)]( + inputs, + weight, + output, + N=shape[1], + K=shape[0], + BLOCK_M=config["block_m"], + BLOCK_N=config["block_n"], + BLOCK_K=config["block_k"], + num_warps=config["warps"], + num_stages=config["stages"], + ) + return output diff --git a/src/sparsevllm/triton_kernel/moe.py b/src/sparsevllm/triton_kernel/moe.py index c8b360e4..cc0e2088 100644 --- a/src/sparsevllm/triton_kernel/moe.py +++ b/src/sparsevllm/triton_kernel/moe.py @@ -8,7 +8,9 @@ from sparsevllm.triton_kernel.silu_and_mul import silu_and_mul_fwd from sparsevllm.triton_kernel.moe_config import ( + MoeGemmConfig, device_info, + resolve_fp8_routed_gemm_config, resolve_moe_gemm_config, ) @@ -89,6 +91,31 @@ def _fill_local_assignments_kernel( tl.store(sorted_token_ids_ptr + positions, assignment_ids, mask=is_local) +@triton.jit +def _prepare_naive_assignment_kernel( + topk_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + num_assignments: tl.constexpr, + local_expert_start: tl.constexpr, + local_expert_end: tl.constexpr, + num_tokens_post_padded: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.arange(0, BLOCK_SIZE) + valid = offsets < num_assignments + global_expert_ids = tl.load(topk_ids_ptr + offsets, mask=valid) + is_local = (global_expert_ids >= local_expert_start) & ( + global_expert_ids < local_expert_end + ) + tl.store( + expert_ids_ptr + offsets, + tl.where(is_local, global_expert_ids - local_expert_start, -1), + mask=valid, + ) + tl.store(num_tokens_post_padded_ptr, num_tokens_post_padded) + + def _validate_alignment_inputs( topk_ids: torch.Tensor, block_size: int, @@ -262,24 +289,25 @@ def _prepare_expert_assignment( ) -> MoeAlignment: num_assignments = int(topk_ids.numel()) if num_assignments * 4 <= int(num_experts): - flat_ids = topk_ids.view(-1) - is_local = (flat_ids >= local_expert_start) & ( - flat_ids < local_expert_end + metadata = torch.empty( + num_assignments + 1, dtype=torch.int32, device=topk_ids.device + ) + expert_ids, num_tokens_post_padded = metadata[:-1], metadata[-1:] + _prepare_naive_assignment_kernel[(1,)]( + topk_ids, + expert_ids, + num_tokens_post_padded, + num_assignments=num_assignments, + local_expert_start=local_expert_start, + local_expert_end=local_expert_end, + num_tokens_post_padded=num_assignments * block_size, + BLOCK_SIZE=triton.next_power_of_2(num_assignments), + num_warps=1, ) - expert_ids = torch.where( - is_local, - flat_ids - local_expert_start, - torch.full_like(flat_ids, -1), - ).to(torch.int32) return MoeAlignment( sorted_token_ids=None, - expert_ids=expert_ids.contiguous(), - num_tokens_post_padded=torch.full( - (1,), - num_assignments * block_size, - dtype=torch.int32, - device=topk_ids.device, - ), + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, block_size=block_size, naive=True, ) @@ -657,6 +685,7 @@ def _routed_fp8_gemm_kernel( INPUT_TOP_K: tl.constexpr, MUL_ROUTING_WEIGHT: tl.constexpr, NAIVE_ASSIGNMENT: tl.constexpr, + SWAP_AB: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, @@ -686,38 +715,70 @@ def _routed_fp8_gemm_kernel( offsets_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) offsets_k = tl.arange(0, BLOCK_SIZE_K) input_rows = assignment_ids // INPUT_TOP_K - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + accumulator = tl.zeros( + (BLOCK_SIZE_N, BLOCK_SIZE_M) if SWAP_AB else (BLOCK_SIZE_M, BLOCK_SIZE_N), + dtype=tl.float32, + ) for k_block in range(0, tl.cdiv(K, BLOCK_SIZE_K)): remaining_k = K - k_block * BLOCK_SIZE_K - a_raw = tl.load( - a_ptr - + input_rows[:, None] * stride_am - + (k_block * BLOCK_SIZE_K + offsets_k[None, :]) * stride_ak, - mask=assignment_mask[:, None] - & (offsets_k[None, :] < remaining_k), - other=0.0, - ).to(tl.float32) - a_scale = tl.max(tl.abs(a_raw), axis=1) / 448.0 - a_quant = (a_raw / tl.maximum(a_scale[:, None], 1.0e-12)).to( - tl.float8e4nv - ) - b = tl.load( - b_ptr - + expert_id * stride_be - + offsets_n[None, :] * stride_bn - + (k_block * BLOCK_SIZE_K + offsets_k[:, None]) * stride_bk, - mask=(offsets_n[None, :] < N) - & (offsets_k[:, None] < remaining_k), - other=0.0, - ) + if SWAP_AB: + a_raw = tl.load( + a_ptr + + input_rows[None, :] * stride_am + + (k_block * BLOCK_SIZE_K + offsets_k[:, None]) * stride_ak, + mask=assignment_mask[None, :] + & (offsets_k[:, None] < remaining_k), + other=0.0, + ).to(tl.float32) + a_scale = tl.max(tl.abs(a_raw), axis=0) / 448.0 + a_quant = (a_raw / tl.maximum(a_scale[None, :], 1.0e-12)).to( + tl.float8e4nv + ) + b = tl.load( + b_ptr + + expert_id * stride_be + + offsets_n[:, None] * stride_bn + + (k_block * BLOCK_SIZE_K + offsets_k[None, :]) * stride_bk, + mask=(offsets_n[:, None] < N) + & (offsets_k[None, :] < remaining_k), + other=0.0, + ) + else: + a_raw = tl.load( + a_ptr + + input_rows[:, None] * stride_am + + (k_block * BLOCK_SIZE_K + offsets_k[None, :]) * stride_ak, + mask=assignment_mask[:, None] + & (offsets_k[None, :] < remaining_k), + other=0.0, + ).to(tl.float32) + a_scale = tl.max(tl.abs(a_raw), axis=1) / 448.0 + a_quant = (a_raw / tl.maximum(a_scale[:, None], 1.0e-12)).to( + tl.float8e4nv + ) + b = tl.load( + b_ptr + + expert_id * stride_be + + offsets_n[None, :] * stride_bn + + (k_block * BLOCK_SIZE_K + offsets_k[:, None]) * stride_bk, + mask=(offsets_n[None, :] < N) + & (offsets_k[:, None] < remaining_k), + other=0.0, + ) b_scale = tl.load( b_scale_ptr + expert_id * stride_bse - + pid_n * stride_bsn + + (pid_n * BLOCK_SIZE_N // 128) * stride_bsn + k_block * stride_bsk ).to(tl.float32) - accumulator += tl.dot(a_quant, b) * a_scale[:, None] * b_scale + if SWAP_AB: + accumulator += tl.dot(b, a_quant) * b_scale * a_scale[None, :] + else: + accumulator += tl.dot(a_quant, b) * a_scale[:, None] * b_scale + + if SWAP_AB: + accumulator = tl.trans(accumulator, (1, 0)) if MUL_ROUTING_WEIGHT: routing_weights = tl.load( @@ -746,7 +807,11 @@ def _routed_fp8_gemm( *, input_top_k: int, multiply_routing_weight: bool, + config: MoeGemmConfig | None = None, ) -> None: + config = config or MoeGemmConfig(alignment.block_size, 128, 128, 1, 4, 3) + if config.block_m != alignment.block_size: + raise ValueError("Routed FP8 GEMM config and assignment block sizes must match.") num_assignments = int(topk_weights.numel()) if alignment.naive: em = num_assignments * alignment.block_size @@ -757,8 +822,8 @@ def _routed_fp8_gemm( em = int(alignment.sorted_token_ids.numel()) sorted_token_ids = alignment.sorted_token_ids grid = ( - triton.cdiv(em, alignment.block_size), - triton.cdiv(int(weights.shape[1]), 128), + triton.cdiv(em, config.block_m), + triton.cdiv(int(weights.shape[1]), config.block_n), ) _routed_fp8_gemm_kernel[grid]( inputs, @@ -786,11 +851,12 @@ def _routed_fp8_gemm( INPUT_TOP_K=int(input_top_k), MUL_ROUTING_WEIGHT=bool(multiply_routing_weight), NAIVE_ASSIGNMENT=alignment.naive, - BLOCK_SIZE_M=alignment.block_size, - BLOCK_SIZE_N=128, - BLOCK_SIZE_K=128, - num_warps=4, - num_stages=3, + SWAP_AB=config.swap_ab, + BLOCK_SIZE_M=config.block_m, + BLOCK_SIZE_N=config.block_n, + BLOCK_SIZE_K=config.block_k, + num_warps=config.num_warps, + num_stages=config.num_stages, ) @@ -1067,7 +1133,6 @@ def fused_moe( local_expert_start=local_expert_start, local_expert_end=local_expert_end, ) - if _fuse_gate_up_swiglu: activated = torch.empty( (num_assignments, intermediate_size), @@ -1274,9 +1339,26 @@ def fused_moe_fp8( num_tokens = int(hidden_states.shape[0]) top_k = int(topk_ids.shape[1]) num_assignments = num_tokens * top_k + device_name, capability = device_info( + hidden_states.device.type, + hidden_states.device.index or 0, + ) + config_kwargs = { + "num_tokens": num_tokens, + "top_k": top_k, + "num_local_experts": num_local_experts, + "hidden_size": hidden_size, + "intermediate_size": intermediate_size, + "device_name": device_name, + "device_capability": capability, + } + w13_config = resolve_fp8_routed_gemm_config(**config_kwargs, stage="w13") + w2_config = resolve_fp8_routed_gemm_config(**config_kwargs, stage="w2") + if w13_config.block_m != w2_config.block_m: + raise ValueError("FP8 routed GEMM stages must share one assignment block size.") alignment = _prepare_expert_assignment( topk_ids, - block_size=16, + block_size=w13_config.block_m, num_experts=num_experts, local_expert_start=local_expert_start, local_expert_end=local_expert_end, @@ -1295,6 +1377,7 @@ def fused_moe_fp8( alignment, input_top_k=top_k, multiply_routing_weight=False, + config=w13_config, ) activated = silu_and_mul_fwd(w13_output, gate_up_order=gate_up_order) w2_output = torch.empty( @@ -1311,6 +1394,7 @@ def fused_moe_fp8( alignment, input_top_k=1, multiply_routing_weight=True, + config=w2_config, ) return moe_sum( w2_output.view(num_tokens, top_k, hidden_size), diff --git a/src/sparsevllm/triton_kernel/moe_config.py b/src/sparsevllm/triton_kernel/moe_config.py index 180e499a..d19837ca 100644 --- a/src/sparsevllm/triton_kernel/moe_config.py +++ b/src/sparsevllm/triton_kernel/moe_config.py @@ -14,6 +14,7 @@ class MoeGemmConfig: group_m: int num_warps: int num_stages: int + swap_ab: bool = False def as_triton_kwargs(self) -> dict[str, int]: return { @@ -92,6 +93,9 @@ def _heuristic_config( _F = MoeGemmConfig(64, 64, 64, 8, 8, 3) _G = MoeGemmConfig(16, 32, 64, 8, 4, 4) _H = MoeGemmConfig(16, 32, 64, 8, 4, 3) +_I = MoeGemmConfig(16, 64, 64, 8, 4, 4) +_J = MoeGemmConfig(16, 64, 64, 8, 4, 2) +_K = MoeGemmConfig(16, 128, 64, 8, 4, 2) def _stage_table( @@ -106,9 +110,21 @@ def _stage_table( } -# Qwen3-30B-A3B BF16 profiles tuned offline on H20. Profiles are keyed by the -# kernel-relevant hardware and GEMM shape rather than by model name. +# BF16 profiles are keyed by kernel-relevant hardware and GEMM shape rather +# than by model name. _TUNED_CONFIGS = { + MoeGemmShape("H20", (9, 0), torch.bfloat16, 8, 256, 2048, 512): { + "w13": {1: _G, 2: _G, 4: _A, 8: _C}, + "w2": {1: _G, 2: _A, 4: _K, 8: _I}, + }, + MoeGemmShape("H20", (9, 0), torch.bfloat16, 8, 256, 2048, 256): { + "w13": {1: _G, 2: _G, 4: _G, 8: _A}, + "w2": {1: _H, 2: _J, 4: _K, 8: _K}, + }, + MoeGemmShape("H20", (9, 0), torch.bfloat16, 8, 128, 2048, 512): { + "w13": {1: _G, 2: _I, 4: _I, 8: _C}, + "w2": {1: _I, 2: _H, 4: _A, 8: _I}, + }, MoeGemmShape("H20", (9, 0), torch.bfloat16, 8, 128, 2048, 768): _stage_table( (_D, _D, _D, _A, _A, _A, _B, _B, _B, _F, _F, _F), (_D, _D, _D, _B, _B, _B, _B, _B, _B, _F, _F, _F), @@ -133,12 +149,78 @@ def _stage_table( (_A, _D, _D, _A, _A, _A, _A, _A, _A, _A, _F, _F), (_A, _D, _D, _A, _A, _A, _A, _A, _A, _A, _A, _B), ), + MoeGemmShape("NVIDIA H100 80GB HBM3", (9, 0), torch.bfloat16, 8, 256, 2048, 512): { + "w13": {1: _G, 2: _I, 4: _A, 8: _A}, + "w2": {1: _G, 2: _G, 4: _A, 8: _H}, + }, + MoeGemmShape("NVIDIA H100 80GB HBM3", (9, 0), torch.bfloat16, 8, 256, 2048, 256): { + "w13": {1: _G, 2: _G, 4: _G, 8: _H}, + "w2": {1: _G, 2: _I, 4: _J, 8: _H}, + }, + MoeGemmShape("NVIDIA H100 80GB HBM3", (9, 0), torch.bfloat16, 8, 128, 2048, 512): { + "w13": {1: _G, 2: _G, 4: _G, 8: _A}, + "w2": {1: _G, 2: _G, 4: _H, 8: _H}, + }, } -# BF16 TP2 expert shards profiled on H100. The fused stage has two FP32 -# accumulators, so reusing the unfused W13 tile creates register pressure. +# The fused BF16 stage has two FP32 accumulators, so reusing a wide unfused W13 +# tile can create register pressure. _TUNED_GATE_UP_SWIGLU_CONFIGS = { + MoeGemmShape( + "H20", + (9, 0), + torch.bfloat16, + 8, + 256, + 2048, + 512, + ): {1: _G, 2: _G, 4: _H, 8: _G}, + MoeGemmShape( + "H20", + (9, 0), + torch.bfloat16, + 8, + 256, + 2048, + 256, + ): {1: _G, 2: _G, 4: _G, 8: _H}, + MoeGemmShape( + "H20", + (9, 0), + torch.bfloat16, + 8, + 128, + 2048, + 512, + ): {1: _G, 2: _G, 4: _G, 8: _H}, + MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.bfloat16, + 8, + 256, + 2048, + 512, + ): {1: _G, 2: _G, 4: _H, 8: _H}, + MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.bfloat16, + 8, + 256, + 2048, + 256, + ): {1: _G, 2: _G, 4: _G, 8: _H}, + MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.bfloat16, + 8, + 128, + 2048, + 512, + ): {1: _G, 2: _G, 4: _G, 8: _H}, MoeGemmShape( "NVIDIA H100 80GB HBM3", (9, 0), @@ -156,6 +238,138 @@ def _stage_table( } +_FP8_N64_SWAP = MoeGemmConfig(16, 64, 128, 1, 4, 3, True) +_FP8_N64_SWAP_S2 = MoeGemmConfig(16, 64, 128, 1, 4, 2, True) +_FP8_N64_SWAP_S4 = MoeGemmConfig(16, 64, 128, 1, 4, 4, True) +_FP8_N64_SWAP_S5 = MoeGemmConfig(16, 64, 128, 1, 4, 5, True) +_FP8_N128 = MoeGemmConfig(16, 128, 128, 1, 4, 3) +_FP8_N128_SWAP_S2 = MoeGemmConfig(16, 128, 128, 1, 4, 2, True) +_FP8_N128_SWAP_S4 = MoeGemmConfig(16, 128, 128, 1, 4, 4, True) + + +# Qwen3.6-35B-A3B block-FP8 decode profiles. Unprofiled token buckets retain +# the explicit generic configuration. +_TUNED_FP8_ROUTED_CONFIGS = { + MoeGemmShape( + "H20", + (9, 0), + torch.float8_e4m3fn, + 8, + 256, + 2048, + 512, + ): { + "w13": { + 1: _FP8_N64_SWAP_S4, + 2: _FP8_N128_SWAP_S4, + 4: _FP8_N64_SWAP_S4, + 8: _FP8_N64_SWAP, + }, + "w2": { + 1: _FP8_N64_SWAP, + 2: _FP8_N64_SWAP_S2, + 4: _FP8_N64_SWAP_S2, + 8: _FP8_N128_SWAP_S2, + }, + }, + MoeGemmShape( + "H20", + (9, 0), + torch.float8_e4m3fn, + 8, + 256, + 2048, + 256, + ): { + "w13": dict.fromkeys((1, 2, 4, 8), _FP8_N64_SWAP_S4), + "w2": { + 1: _FP8_N64_SWAP, + 2: _FP8_N64_SWAP_S2, + 4: _FP8_N64_SWAP_S2, + 8: _FP8_N128_SWAP_S2, + }, + }, + MoeGemmShape( + "H20", + (9, 0), + torch.float8_e4m3fn, + 8, + 128, + 2048, + 512, + ): { + "w13": dict.fromkeys((1, 2, 4, 8), _FP8_N64_SWAP_S4), + "w2": { + 1: _FP8_N64_SWAP, + 2: _FP8_N128_SWAP_S4, + 4: _FP8_N64_SWAP_S2, + 8: _FP8_N64_SWAP_S2, + }, + }, + MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.float8_e4m3fn, + 8, + 256, + 2048, + 512, + ): { + "w13": { + 1: _FP8_N64_SWAP_S5, + 2: _FP8_N64_SWAP_S4, + 4: _FP8_N64_SWAP, + 8: _FP8_N64_SWAP_S4, + }, + "w2": { + 1: _FP8_N64_SWAP_S4, + 2: _FP8_N128_SWAP_S4, + 4: _FP8_N64_SWAP_S2, + 8: _FP8_N64_SWAP, + }, + }, + MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.float8_e4m3fn, + 8, + 256, + 2048, + 256, + ): { + "w13": { + 1: _FP8_N64_SWAP_S4, + 2: _FP8_N64_SWAP_S5, + 4: _FP8_N64_SWAP_S4, + 8: _FP8_N64_SWAP, + }, + "w2": { + 1: _FP8_N64_SWAP, + 2: _FP8_N64_SWAP_S2, + 4: _FP8_N64_SWAP_S2, + 8: _FP8_N64_SWAP_S2, + }, + }, + MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.float8_e4m3fn, + 8, + 128, + 2048, + 512, + ): { + "w13": dict.fromkeys((1, 2, 4, 8), _FP8_N64_SWAP_S4), + "w2": { + 1: _FP8_N64_SWAP_S4, + 2: _FP8_N64_SWAP, + 4: _FP8_N64_SWAP_S2, + 8: _FP8_N64_SWAP_S2, + }, + }, +} + + @lru_cache(maxsize=None) def _resolve_moe_gemm_config( dtype: torch.dtype, @@ -181,7 +395,9 @@ def _resolve_moe_gemm_config( if stage == "gate_up_swiglu": fused_table = _TUNED_GATE_UP_SWIGLU_CONFIGS.get(shape) if fused_table is not None: - return fused_table[token_bucket(num_tokens)] + tuned = fused_table.get(token_bucket(num_tokens)) + if tuned is not None: + return tuned assignments = num_tokens * top_k return MoeGemmConfig( block_m=16, @@ -192,7 +408,9 @@ def _resolve_moe_gemm_config( num_stages=4 if assignments <= 64 else 3, ) if table is not None: - return table[stage][token_bucket(num_tokens)] + tuned = table[stage].get(token_bucket(num_tokens)) + if tuned is not None: + return tuned output_size = 2 * intermediate_size if stage == "w13" else hidden_size return _heuristic_config( @@ -234,3 +452,35 @@ def resolve_moe_gemm_config( device_name, device_capability, ) + + +def resolve_fp8_routed_gemm_config( + *, + num_tokens: int, + top_k: int, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + stage: str, + device_name: str | None = None, + device_capability: tuple[int, int] | None = None, +) -> MoeGemmConfig: + if stage not in {"w13", "w2"}: + raise ValueError(f"FP8 routed GEMM stage must be 'w13' or 'w2', got {stage!r}.") + if device_name is None: + device_name = torch.cuda.get_device_name() + if device_capability is None: + device_capability = torch.cuda.get_device_capability() + shape = MoeGemmShape( + _hardware_family(device_name), + device_capability, + torch.float8_e4m3fn, + int(top_k), + int(num_local_experts), + int(hidden_size), + int(intermediate_size), + ) + tuned = _TUNED_FP8_ROUTED_CONFIGS.get(shape, {}).get(stage, {}).get( + token_bucket(num_tokens) + ) + return tuned or _FP8_N128 diff --git a/src/sparsevllm/triton_kernel/moe_topk.py b/src/sparsevllm/triton_kernel/moe_topk.py index 5c115d67..e9d1f821 100644 --- a/src/sparsevllm/triton_kernel/moe_topk.py +++ b/src/sparsevllm/triton_kernel/moe_topk.py @@ -76,6 +76,7 @@ def _gather_candidate( equal_rank, num_greater, slot: tl.constexpr, + num_experts: tl.constexpr, ): use_greater = slot < num_greater rank = tl.where(use_greater, slot, slot - num_greater) @@ -84,7 +85,7 @@ def _gather_candidate( greater_mask & (greater_rank == rank), equal_mask & (equal_rank == rank), ) - expert_id = tl.min(tl.where(mask, offsets, 128), axis=0) + expert_id = tl.min(tl.where(mask, offsets, num_experts), axis=0) value = tl.sum(tl.where(mask, values, 0.0), axis=0) return value, expert_id @@ -97,9 +98,10 @@ def _topk_softmax_kernel( stride_logits_m, stride_weights_m, stride_ids_m, + NUM_EXPERTS: tl.constexpr, ): row = tl.program_id(0) - offsets = tl.arange(0, 128) + offsets = tl.arange(0, NUM_EXPERTS) logits = tl.load(logits_ptr + row * stride_logits_m + offsets).to(tl.float32) row_max = tl.max(logits, axis=0) probabilities = libdevice.exp(logits - row_max) @@ -119,35 +121,35 @@ def _topk_softmax_kernel( v0, i0 = _gather_candidate( probabilities, offsets, greater_mask, equal_mask, - greater_rank, equal_rank, num_greater, 0, + greater_rank, equal_rank, num_greater, 0, NUM_EXPERTS, ) v1, i1 = _gather_candidate( probabilities, offsets, greater_mask, equal_mask, - greater_rank, equal_rank, num_greater, 1, + greater_rank, equal_rank, num_greater, 1, NUM_EXPERTS, ) v2, i2 = _gather_candidate( probabilities, offsets, greater_mask, equal_mask, - greater_rank, equal_rank, num_greater, 2, + greater_rank, equal_rank, num_greater, 2, NUM_EXPERTS, ) v3, i3 = _gather_candidate( probabilities, offsets, greater_mask, equal_mask, - greater_rank, equal_rank, num_greater, 3, + greater_rank, equal_rank, num_greater, 3, NUM_EXPERTS, ) v4, i4 = _gather_candidate( probabilities, offsets, greater_mask, equal_mask, - greater_rank, equal_rank, num_greater, 4, + greater_rank, equal_rank, num_greater, 4, NUM_EXPERTS, ) v5, i5 = _gather_candidate( probabilities, offsets, greater_mask, equal_mask, - greater_rank, equal_rank, num_greater, 5, + greater_rank, equal_rank, num_greater, 5, NUM_EXPERTS, ) v6, i6 = _gather_candidate( probabilities, offsets, greater_mask, equal_mask, - greater_rank, equal_rank, num_greater, 6, + greater_rank, equal_rank, num_greater, 6, NUM_EXPERTS, ) v7, i7 = _gather_candidate( probabilities, offsets, greater_mask, equal_mask, - greater_rank, equal_rank, num_greater, 7, + greater_rank, equal_rank, num_greater, 7, NUM_EXPERTS, ) values_and_ids = _torch_sort8( v0, i0, v1, i1, v2, i2, v3, i3, @@ -196,13 +198,15 @@ def topk_softmax( "Triton topk_softmax supports BF16 and FP16 logits, got " f"{router_logits.dtype}." ) - if not router_logits.is_contiguous(): - raise ValueError("Triton topk_softmax requires contiguous router_logits.") + if router_logits.stride(1) != 1: + raise ValueError("Triton topk_softmax requires a contiguous expert dimension.") if int(router_logits.shape[0]) <= 0: raise ValueError("Triton topk_softmax requires at least one token.") - if int(router_logits.shape[1]) != 128 or int(top_k) != 8: + num_experts = int(router_logits.shape[1]) + if num_experts not in {128, 256} or int(top_k) != 8: raise ValueError( - "Triton topk_softmax currently requires num_experts=128 and top_k=8, " + "Triton topk_softmax currently requires num_experts=128 or " + "num_experts=256 and top_k=8, " f"got num_experts={router_logits.shape[1]}, top_k={top_k}." ) @@ -223,6 +227,7 @@ def topk_softmax( router_logits.stride(0), weights.stride(0), ids.stride(0), + NUM_EXPERTS=num_experts, num_warps=1, ) return weights, ids diff --git a/src/sparsevllm/triton_kernel/qwen3_5/gated_shared_add.py b/src/sparsevllm/triton_kernel/qwen3_5/gated_shared_add.py new file mode 100644 index 00000000..1fe35f8e --- /dev/null +++ b/src/sparsevllm/triton_kernel/qwen3_5/gated_shared_add.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl +import triton.language.extra.libdevice as libdevice + + +@triton.jit +def _gated_shared_add_kernel( + routed_ptr, + shared_ptr, + gate_ptr, + output_ptr, + gate_stride, + hidden_size: tl.constexpr, + block_size: tl.constexpr, +): + row, block = tl.program_id(0), tl.program_id(1) + offsets = block * block_size + tl.arange(0, block_size) + mask = offsets < hidden_size + routed = tl.load(routed_ptr + row * hidden_size + offsets, mask=mask) + shared = tl.load(shared_ptr + row * hidden_size + offsets, mask=mask) + gate = tl.load(gate_ptr + row * gate_stride).to(tl.float32) + gate = (1.0 / (1.0 + libdevice.exp(-gate))).to(tl.bfloat16) + tl.store( + output_ptr + row * hidden_size + offsets, + routed + gate * shared, + mask=mask, + ) + + +def triton_gated_shared_add( + routed: torch.Tensor, + shared: torch.Tensor, + gate_logits: torch.Tensor, +) -> torch.Tensor: + output = torch.empty_like(routed) + hidden_size = int(routed.shape[1]) + block_size = 512 + _gated_shared_add_kernel[(routed.shape[0], triton.cdiv(hidden_size, block_size))]( + routed, + shared, + gate_logits, + output, + gate_logits.stride(0), + hidden_size=hidden_size, + block_size=block_size, + num_warps=4, + ) + return output diff --git a/src/sparsevllm/utils/config.py b/src/sparsevllm/utils/config.py new file mode 100644 index 00000000..665270c9 --- /dev/null +++ b/src/sparsevllm/utils/config.py @@ -0,0 +1,7 @@ +from typing import Any + + +def config_get(config: Any, name: str, default: Any = None) -> Any: + if config is None: + return default + return config.get(name, default) if isinstance(config, dict) else getattr(config, name, default) diff --git a/src/sparsevllm/utils/loader.py b/src/sparsevllm/utils/loader.py index 67ae5ef8..b6b47d86 100644 --- a/src/sparsevllm/utils/loader.py +++ b/src/sparsevllm/utils/loader.py @@ -19,6 +19,16 @@ def default_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor): param.data.copy_(loaded_weight) +def get_parameter_or_buffer(model: nn.Module, name: str) -> torch.Tensor: + try: + return model.get_parameter(name) + except AttributeError as parameter_error: + try: + return model.get_buffer(name) + except AttributeError: + raise parameter_error + + @dataclass(frozen=True) class TensorMetadata: shape: tuple[int, ...] @@ -70,12 +80,10 @@ def _rank_local_slice_for_tensor( source_weight_name: str, source_shape: tuple[int, ...], ) -> tuple[slice, ...] | None: - is_scale = source_weight_name.endswith(".weight_scale_inv") - source_parameter_name = ( - source_weight_name[: -len(".weight_scale_inv")] + ".weight" - if is_scale - else source_weight_name - ) + source_parameter_name = _weight_key_for_scale_key(model, source_weight_name) + is_scale = source_parameter_name is not None + if source_parameter_name is None: + source_parameter_name = source_weight_name target_parameter_name = _target_weight_name_for_model( model, source_parameter_name ) @@ -112,12 +120,7 @@ def _read_safetensors_shard( tensors[key] = handle.get_tensor(key) continue - is_scale = key.endswith(".weight_scale_inv") - source_parameter_name = ( - key[: -len(".weight_scale_inv")] + ".weight" - if is_scale - else key - ) + source_parameter_name = _weight_key_for_scale_key(model, key) or key if _target_weight_name_for_model(model, source_parameter_name) is None: continue rank_slice = _rank_local_slice_for_tensor( @@ -175,12 +178,38 @@ def _module_for_parameter(model: nn.Module, param_name: str) -> nn.Module: return model.get_submodule(module_name) -def _scale_key_for_weight_key(weight_key: str) -> str: +def _scale_key_for_weight_key(model: nn.Module, weight_key: str) -> str: if not weight_key.endswith(".weight"): raise ValueError(f"Expected a weight key ending in '.weight', got {weight_key!r}.") + resolver = getattr(model, "scale_key_for_weight", None) + if callable(resolver): + scale_key = resolver(weight_key) + if scale_key is not None: + if not isinstance(scale_key, str): + raise TypeError( + "scale_key_for_weight() must return str or None, " + f"got {type(scale_key).__name__}." + ) + return scale_key return weight_key[: -len(".weight")] + ".weight_scale_inv" +def _weight_key_for_scale_key(model: nn.Module, scale_key: str) -> str | None: + resolver = getattr(model, "weight_key_for_scale", None) + if callable(resolver): + weight_key = resolver(scale_key) + if weight_key is not None and not isinstance(weight_key, str): + raise TypeError( + "weight_key_for_scale() must return str or None, " + f"got {type(weight_key).__name__}." + ) + if weight_key is not None: + return weight_key + if scale_key.endswith(".weight_scale_inv"): + return scale_key[: -len(".weight_scale_inv")] + ".weight" + return None + + def _target_weight_name_for_model(model: nn.Module, source_weight_name: str) -> str | None: ignored_prefixes = tuple(getattr(model, "ignored_weight_prefixes", ())) if source_weight_name.startswith(ignored_prefixes): @@ -696,14 +725,18 @@ def load_model( f"{duplicate_source_keys[:5]}." ) seen_source_keys.update(keys) - scale_keys = {key for key in keys if key.endswith(".weight_scale_inv")} + scale_keys = { + _scale_key_for_weight_key(model, key) + for key in keys + if key.endswith(".weight") + }.intersection(keys) consumed_scale_keys: set[str] = set() for source_weight_name in keys: - if source_weight_name.endswith(".weight_scale_inv"): + if source_weight_name in scale_keys: continue scale_key = None if source_weight_name.endswith(".weight"): - scale_key = _scale_key_for_weight_key(source_weight_name) + scale_key = _scale_key_for_weight_key(model, source_weight_name) param_name = _target_weight_name_for_model(model, source_weight_name) if param_name is None: skipped_weight_hook = getattr(model, "record_skipped_weight", None) @@ -778,7 +811,7 @@ def load_model( module = _module_for_parameter(model, param_name) loaded_scale = None if source_weight_name.endswith(".weight"): - scale_key = _scale_key_for_weight_key(source_weight_name) + scale_key = _scale_key_for_weight_key(model, source_weight_name) loaded_scale = tensors.get(scale_key) if loaded_scale is not None: consumed_scale_keys.add(scale_key) @@ -793,7 +826,7 @@ def load_model( loaded_parameter_names.add(param_name) loaded_count += 1 continue - param = model.get_parameter(param_name) + param = get_parameter_or_buffer(model, param_name) weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, tensors[source_weight_name]) loaded_parameter_names.add(param_name) diff --git a/src/sparsevllm/utils/select_omnikv_full_layers.py b/src/sparsevllm/utils/select_omnikv_full_layers.py index 638961cb..435b151d 100644 --- a/src/sparsevllm/utils/select_omnikv_full_layers.py +++ b/src/sparsevllm/utils/select_omnikv_full_layers.py @@ -71,19 +71,12 @@ def text_model_config(config): def attention_layer_indices_from_config(config) -> list[int]: + from sparsevllm.models.layout import RuntimeLayout + text_config = text_model_config(config) - num_hidden_layers = int(text_config.num_hidden_layers) - layer_types = getattr(text_config, "layer_types", None) - if layer_types is None: - return list(range(num_hidden_layers)) - if len(layer_types) != num_hidden_layers: - raise ValueError( - "layer_types length does not match num_hidden_layers: " - f"{len(layer_types)} != {num_hidden_layers}." - ) - indices = [idx for idx, layer_type in enumerate(layer_types) if layer_type == "full_attention"] + indices = list(RuntimeLayout.from_config(text_config).full_attention_layer_indices) if not indices: - raise ValueError("Model layer_types does not contain any full_attention layer.") + raise ValueError("Model layout does not contain any full-attention layer.") return indices diff --git a/tests/test_deepseek_v4_config.py b/tests/test_deepseek_v4_config.py new file mode 100644 index 00000000..40714baa --- /dev/null +++ b/tests/test_deepseek_v4_config.py @@ -0,0 +1,162 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from sparsevllm.config import Config +from sparsevllm.debug.tiny_random import apply_tiny_random_overrides +from sparsevllm.distributed import ParallelMode +from sparsevllm.method_registry import ( + DEEPSEEK_V4_DPA_EP_COMPATIBILITY, + MODEL_RUNTIME_COMPATIBILITY, +) + + +def _official_config(**overrides): + values = { + "architectures": ["DeepseekV4ForCausalLM"], + "model_type": "deepseek_v4", + "vocab_size": 129280, + "hidden_size": 4096, + "intermediate_size": 2048, + "moe_intermediate_size": 2048, + "num_hidden_layers": 43, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "head_dim": 512, + "q_lora_rank": 1024, + "o_lora_rank": 1024, + "o_groups": 8, + "qk_rope_head_dim": 64, + "index_n_heads": 64, + "index_head_dim": 128, + "index_topk": 512, + "n_routed_experts": 256, + "num_local_experts": 256, + "num_experts_per_tok": 6, + "expert_dtype": "fp4", + "hc_mult": 4, + "sliding_window": 128, + "compress_rates": { + "compressed_sparse_attention": 4, + "heavily_compressed_attention": 128, + }, + "num_nextn_predict_layers": 1, + "max_position_embeddings": 1048576, + "torch_dtype": torch.bfloat16, + "layer_types": ["sliding_attention", "sliding_attention"] + + ["compressed_sparse_attention", "heavily_compressed_attention"] * 20 + + ["compressed_sparse_attention"], + "mlp_layer_types": ["hash_moe"] * 3 + ["moe"] * 40, + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + }, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _make_config(tmp_path, hf_config=None, **kwargs): + hf_config = hf_config or _official_config() + with patch( + "sparsevllm.configs.runtime.AutoConfig.from_pretrained", + return_value=hf_config, + ): + return Config(model=str(tmp_path), **kwargs) + + +@pytest.mark.parametrize("parallel_size", [1, 2, 4]) +def test_deepseek_v4_accepts_overlapping_dpa_ep(tmp_path, parallel_size): + config = _make_config( + tmp_path, + data_parallel_size=parallel_size, + expert_parallel_size=parallel_size, + decode_cuda_graph=True, + enforce_eager=False, + ) + + assert config.uses_dpa_ep_layout + assert config.world_size == parallel_size + assert config.attention_tensor_parallel_size == 1 + assert config.moe_tensor_parallel_size == 1 + assert config.parallel_topology.mode is ParallelMode.DPA_EP + assert ( + MODEL_RUNTIME_COMPATIBILITY[("deepseek_v4", ParallelMode.DPA_EP)] + is DEEPSEEK_V4_DPA_EP_COMPATIBILITY + ) + + +@pytest.mark.parametrize( + "parallel_kwargs", + [ + {"tensor_parallel_size": 2}, + {"data_parallel_size": 2, "expert_parallel_size": 1}, + {"data_parallel_size": 1, "expert_parallel_size": 2}, + {"data_parallel_size": 3, "expert_parallel_size": 3}, + ], +) +def test_deepseek_v4_rejects_invalid_parallel_layout(tmp_path, parallel_kwargs): + with pytest.raises(ValueError, match=r"DPA\+EP|DeepSeek V4"): + _make_config(tmp_path, **parallel_kwargs) + + +def test_deepseek_v4_rejects_sparse_methods_and_prefix_cache(tmp_path): + with pytest.raises(ValueError, match="validated methods"): + _make_config(tmp_path, vllm_sparse_method="quest") + with pytest.raises(ValueError, match="prefix caching"): + _make_config(tmp_path, enable_prefix_caching=True) + + +@pytest.mark.parametrize( + ("field_name", "invalid_value"), + [ + ("head_dim", 128), + ("index_topk", 256), + ("n_routed_experts", 128), + ("hc_mult", 2), + ("num_nextn_predict_layers", 0), + ("torch_dtype", torch.float16), + ("expert_dtype", "bf16"), + ], +) +def test_deepseek_v4_rejects_checkpoint_drift(tmp_path, field_name, invalid_value): + expected_field = field_name.replace("n_routed", "routed").replace("torch_", "") + with pytest.raises(ValueError, match=expected_field): + _make_config(tmp_path, hf_config=_official_config(**{field_name: invalid_value})) + + +def test_deepseek_v4_tiny_random_shrinks_lists_and_disables_quantization(tmp_path): + tiny_path = tmp_path / "tiny.json" + tiny_path.write_text( + '{"num_hidden_layers": 4, "hidden_size": 128, "intermediate_size": 64, ' + '"head_dim": 64, "num_attention_heads": 4, "q_lora_rank": 64, ' + '"o_lora_rank": 64, "o_groups": 2, "index_n_heads": 4, ' + '"index_head_dim": 64, "index_topk": 8, "num_local_experts": 8, ' + '"n_routed_experts": 8, "num_experts_per_tok": 2, "sliding_window": 8, ' + '"vocab_size": 128}', + encoding="utf-8", + ) + hf_config = _official_config() + applied = apply_tiny_random_overrides(hf_config, str(tiny_path)) + + assert applied["num_hidden_layers"] == 4 + assert hf_config.layer_types == [ + "sliding_attention", + "sliding_attention", + "compressed_sparse_attention", + "heavily_compressed_attention", + ] + assert hf_config.mlp_layer_types == ["hash_moe", "hash_moe", "hash_moe", "moe"] + + config = _make_config( + tmp_path, + hf_config=_official_config(), + tiny_random=True, + tiny_random_config=str(tiny_path), + ) + assert not config.quantization_config.enabled diff --git a/tests/test_deepseek_v4_model.py b/tests/test_deepseek_v4_model.py new file mode 100644 index 00000000..5e1f9e17 --- /dev/null +++ b/tests/test_deepseek_v4_model.py @@ -0,0 +1,359 @@ +import copy +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import pytest +from transformers import DeepseekV4Config, DynamicCache +from transformers.masking_utils import create_sliding_window_causal_mask +from transformers.models.deepseek_v4.modeling_deepseek_v4 import ( + DeepseekV4Attention as ReferenceDeepseekV4Attention, + DeepseekV4RotaryEmbedding, +) + +from sparsevllm.models.deepseek_v4 import DeepseekV4ForCausalLM +from sparsevllm.engine.model_runner import ModelRunner +from sparsevllm.models.deepseek_v4_native import ( + DeepseekV4GroupedFp8Linear, + DeepseekV4HyperConnection, + DeepseekV4PackedExperts, + DeepseekV4Attention, +) +from sparsevllm.operators.moe import FlashInferCutlassFp4MoeProvider +from sparsevllm.utils.context import reset_context, set_context + + +def _config(): + config = DeepseekV4Config( + vocab_size=32, + hidden_size=64, + moe_intermediate_size=32, + num_hidden_layers=4, + num_attention_heads=2, + head_dim=32, + q_lora_rank=32, + num_experts_per_tok=2, + n_routed_experts=4, + n_shared_experts=1, + layer_types=[ + "sliding_attention", + "sliding_attention", + "compressed_sparse_attention", + "heavily_compressed_attention", + ], + mlp_layer_types=["hash_moe", "hash_moe", "hash_moe", "moe"], + sliding_window=8, + o_groups=2, + o_lora_rank=32, + index_n_heads=2, + index_head_dim=32, + index_topk=4, + partial_rotary_factor=1.0, + dtype=torch.float32, + ) + config.sparsevllm_tiny_random = True + return config + + +def _parallel_context(): + return SimpleNamespace( + tp_rank=0, + tp_size=1, + tp_all_reduce=lambda tensor: tensor, + tp_gather=lambda tensor: [tensor], + ) + + +def test_tiny_reference_prefill_and_decode_match_full_sequence(): + with patch( + "sparsevllm.layers.embed_head.get_parallel_context", + return_value=_parallel_context(), + ): + torch.manual_seed(7) + model = DeepseekV4ForCausalLM(_config()).eval() + reference = copy.deepcopy(model.model).eval() + input_ids = torch.arange(9, dtype=torch.long).remainder(model.config.vocab_size) + positions = torch.arange(9, dtype=torch.long) + seq = SimpleNamespace(seq_id=11) + + try: + set_context( + True, + cu_seqlens_q=torch.tensor([0, input_ids.numel()]), + seqs=[seq], + ) + with torch.inference_mode(): + actual_prefill = model(input_ids, positions) + expected_prefill = reference( + input_ids=input_ids.unsqueeze(0), + position_ids=positions.unsqueeze(0), + use_cache=False, + return_dict=True, + ).last_hidden_state.squeeze(0) + torch.testing.assert_close(actual_prefill, expected_prefill) + + set_context(False, seqs=[seq]) + with torch.inference_mode(): + actual_decode = model(torch.tensor([9]), torch.tensor([9])) + expected_decode = reference( + input_ids=torch.arange(10).unsqueeze(0), + position_ids=torch.arange(10).unsqueeze(0), + use_cache=False, + return_dict=True, + ).last_hidden_state[:, -1] + torch.testing.assert_close(actual_decode, expected_decode, rtol=1e-5, atol=2e-6) + finally: + reset_context() + + +def test_formal_runtime_selects_native_model(): + config = _config() + config.sparsevllm_tiny_random = False + native = torch.nn.Module() + with ( + patch( + "sparsevllm.models.deepseek_v4.NativeDeepseekV4Model", + return_value=native, + ), + patch( + "sparsevllm.layers.embed_head.get_parallel_context", + return_value=_parallel_context(), + ), + ): + model = DeepseekV4ForCausalLM(config) + + assert model.model is native + assert not model.tiny_random + + +def test_native_hyper_connection_matches_transformers_reference(): + from transformers.models.deepseek_v4.modeling_deepseek_v4 import ( + DeepseekV4HyperConnection as ReferenceHyperConnection, + ) + + config = _config() + torch.manual_seed(9) + reference = ReferenceHyperConnection(config) + torch.nn.init.normal_(reference.fn, std=0.02) + torch.nn.init.zeros_(reference.base) + torch.nn.init.ones_(reference.scale) + native = DeepseekV4HyperConnection(config) + native.load_state_dict(reference.state_dict()) + streams = torch.randn(2, 5, config.hc_mult, config.hidden_size) + + actual = native(streams) + expected = reference(streams) + + for actual_tensor, expected_tensor in zip(actual, expected): + torch.testing.assert_close(actual_tensor, expected_tensor) + + +def test_native_grouped_linear_preserves_group_boundaries(): + config = _config() + config.quantization_config = SimpleNamespace(enabled=False) + parallel = _parallel_context() + with patch( + "sparsevllm.layers.linear.get_parallel_context", + return_value=parallel, + ): + grouped = DeepseekV4GroupedFp8Linear(config) + torch.manual_seed(4) + grouped.weight.data.normal_() + x = torch.randn(3, config.o_groups, config.num_attention_heads * config.head_dim // config.o_groups) + weight = grouped.weight.view(config.o_groups, config.o_lora_rank, -1) + + actual = grouped(x) + expected = torch.einsum("tgi,goi->tgo", x, weight) + + torch.testing.assert_close(actual, expected) + + +def test_native_fp4_expert_loader_preserves_checkpoint_bits(): + config = SimpleNamespace( + n_routed_experts=4, + num_experts_per_tok=2, + hidden_size=256, + moe_intermediate_size=128, + dtype=torch.bfloat16, + decode_cuda_graph=True, + swiglu_limit=10.0, + ) + parallel = SimpleNamespace(ep_rank=1, ep_size=2) + with ( + patch( + "sparsevllm.models.deepseek_v4_native.get_parallel_context", + return_value=parallel, + ), + patch( + "sparsevllm.models.deepseek_v4_native.resolve_moe_provider", + return_value=FlashInferCutlassFp4MoeProvider(), + ), + ): + experts = DeepseekV4PackedExperts(config) + weight_bits = torch.arange(128 * 128, dtype=torch.int64).to(torch.uint8).view(128, 128) + scale_bits = torch.arange(128 * 8, dtype=torch.int64).to(torch.uint8).view(128, 8) + + experts.load_expert_weight( + 2, + "gate", + weight_bits.view(torch.int8), + scale_bits.view(torch.float8_e8m0fnu), + ) + + assert torch.equal(experts.w13_weight[0, 128:], weight_bits) + assert torch.equal(experts.w13_scale_inv[0, 128:], scale_bits) + + +class _AttentionCache: + def __init__(self, config): + rows, max_len = 1, 256 + head_dim, index_dim = int(config.head_dim), int(config.index_head_dim) + dtype = config.torch_dtype + self.sliding_window = int(config.sliding_window) + self.max_model_len = max_len + self._decode_static_max_context_len = None + self.raw_kv = torch.zeros(1, rows, self.sliding_window, head_dim, dtype=dtype) + self.csa_kv = torch.zeros(1, rows, max_len // 4, head_dim, dtype=dtype) + self.csa_index = torch.zeros(1, rows, max_len // 4, index_dim, dtype=dtype) + self.hca_kv = torch.zeros(1, rows, max_len // 128, head_dim, dtype=dtype) + self.csa_ring_kv = torch.zeros(1, rows, 4, 2 * head_dim, dtype=dtype) + self.csa_ring_gate = torch.full_like(self.csa_ring_kv, float("-inf")) + self.csa_overlap_kv = torch.zeros(1, rows, 4, head_dim, dtype=dtype) + self.csa_overlap_gate = torch.full_like(self.csa_overlap_kv, float("-inf")) + self.index_ring_kv = torch.zeros(1, rows, 4, 2 * index_dim, dtype=dtype) + self.index_ring_gate = torch.full_like(self.index_ring_kv, float("-inf")) + self.index_overlap_kv = torch.zeros(1, rows, 4, index_dim, dtype=dtype) + self.index_overlap_gate = torch.full_like(self.index_overlap_kv, float("-inf")) + self.hca_ring_kv = torch.zeros(1, rows, 128, head_dim, dtype=dtype) + self.hca_ring_gate = torch.full_like(self.hca_ring_kv, float("-inf")) + self.state = SimpleNamespace(req_indices=torch.tensor([0])) + + def csa_slot(self, layer_idx): + return 0 + + def hca_slot(self, layer_idx): + return 0 + + def compressed_capacity(self, ratio, positions): + return (int(positions.max()) + int(ratio)) // int(ratio) + + def get_layer_batch_states(self, layer_idx): + return self.state + + +@pytest.mark.parametrize( + ("layer_type", "prefill_tokens"), + [ + ("sliding_attention", 13), + ("compressed_sparse_attention", 13), + ("heavily_compressed_attention", 140), + ], +) +def test_native_attention_prefill_and_decode_match_reference(layer_type, prefill_tokens): + config = _config() + config.num_hidden_layers = 1 + config.layer_types = [layer_type] + config.mlp_layer_types = ["hash_moe"] + config.quantization_config = SimpleNamespace(enabled=False) + config._attn_implementation = "eager" + parallel = _parallel_context() + with patch( + "sparsevllm.layers.linear.get_parallel_context", + return_value=parallel, + ): + native = DeepseekV4Attention(config, 0).eval() + reference = ReferenceDeepseekV4Attention(config, 0).eval() + torch.manual_seed(29) + for parameter in reference.parameters(): + torch.nn.init.normal_(parameter, std=0.02) + native.load_state_dict(reference.state_dict()) + rotary = DeepseekV4RotaryEmbedding(config) + manager = _AttentionCache(config) + cache = DynamicCache(config=config) + torch.manual_seed(31) + hidden_states = torch.randn(1, prefill_tokens, config.hidden_size) + positions = torch.arange(prefill_tokens).unsqueeze(0) + embeddings = { + kind: rotary(hidden_states, positions, kind) + for kind in ("main", "compress") + } + mask = create_sliding_window_causal_mask( + config, hidden_states, None, cache, positions + ) + + with torch.inference_mode(): + expected = reference(hidden_states, embeddings, positions, mask, cache)[0] + try: + set_context(True, cache_manager=manager, seqs=[SimpleNamespace(seq_id=1)]) + with torch.inference_mode(): + actual = native( + hidden_states, + embeddings, + positions, + cache_rows=manager.state.req_indices, + )[0] + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=2e-6) + + hidden_states = torch.randn(1, 1, config.hidden_size) + positions = torch.tensor([[prefill_tokens]]) + embeddings = { + kind: rotary(hidden_states, positions, kind) + for kind in ("main", "compress") + } + with torch.inference_mode(): + expected = reference(hidden_states, embeddings, positions, None, cache)[0] + set_context(False, cache_manager=manager, seqs=[SimpleNamespace(seq_id=1)]) + with torch.inference_mode(): + actual = native( + hidden_states, + embeddings, + positions, + cache_rows=manager.state.req_indices, + )[0] + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=2e-6) + finally: + reset_context() + + +def test_checkpoint_maps_hyper_head_parameters(): + config = _config() + config.sparsevllm_tiny_random = False + native = torch.nn.Module() + with ( + patch("sparsevllm.models.deepseek_v4.NativeDeepseekV4Model", return_value=native), + patch("sparsevllm.layers.embed_head.get_parallel_context", return_value=_parallel_context()), + ): + model = DeepseekV4ForCausalLM(config) + + assert model.map_weight_name("hc_head_fn") == "model.hc_head.fn" + assert model.map_weight_name("hc_head_base") == "model.hc_head.base" + assert model.map_weight_name("hc_head_scale") == "model.hc_head.scale" + + +def test_dpa_decode_partition_balances_stable_sequence_owners(): + seqs = [SimpleNamespace(seq_id=seq_id) for seq_id in (0, 4, 8)] + for rank in range(4): + runner = ModelRunner.__new__(ModelRunner) + runner.parallel_context = SimpleNamespace(dp_size=4, dp_rank=rank) + selected, owned = runner._deepseek_v4_decode_partition(seqs) + + assert len(selected) == 3 + assert owned == (3 if rank == 0 else 0) + assert all(int(seq.seq_id) % 4 == rank for seq in selected[:owned]) + + +def test_dpa_logits_are_restored_to_global_sequence_order(): + runner = ModelRunner.__new__(ModelRunner) + gathered = torch.tensor([[10.0], [40.0], [20.0], [0.0]]) + runner.parallel_context = SimpleNamespace( + dp_size=2, + ep_all_gather_into_tensor=lambda tensor: gathered, + ) + runner.rank = 0 + seqs = [SimpleNamespace(seq_id=0), SimpleNamespace(seq_id=1), SimpleNamespace(seq_id=4)] + + actual = runner._gather_deepseek_v4_logits( + torch.tensor([[10.0], [40.0]]), seqs, local_owned=2 + ) + + torch.testing.assert_close(actual, torch.tensor([[10.0], [20.0], [40.0]])) diff --git a/tests/test_minimax_m2_config.py b/tests/test_minimax_m2_config.py index 61f75722..a9e32c78 100644 --- a/tests/test_minimax_m2_config.py +++ b/tests/test_minimax_m2_config.py @@ -5,6 +5,7 @@ import torch from sparsevllm.config import Config +from sparsevllm.distributed import ParallelMode, ParallelTopology from sparsevllm.method_registry import ( MINIMAX_M2_EP_COMPATIBILITY, MINIMAX_M2_TP_EP_COMPATIBILITY, @@ -92,15 +93,12 @@ def test_minimax_config_accepts_first_milestone_runtime( @pytest.mark.parametrize( ("field_name", "invalid_value"), [ - ("hidden_size", 4096), - ("rotary_dim", 128), ("qk_norm_type", "per_head"), ("scoring_func", "softmax"), ("use_mtp", False), - ("num_mtp_modules", 0), ], ) -def test_minimax_config_rejects_checkpoint_drift( +def test_minimax_config_rejects_unsupported_semantics( tmp_path, field_name, invalid_value, @@ -110,6 +108,30 @@ def test_minimax_config_rejects_checkpoint_drift( _make_config(tmp_path, hf_config=hf_config) +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("vocab_size", 200192), + ("hidden_size", 4096), + ("num_hidden_layers", 48), + ("rotary_dim", 128), + ("max_position_embeddings", 262144), + ("num_mtp_modules", 0), + ], +) +def test_minimax_config_accepts_checkpoint_shape_variants( + tmp_path, + field_name, + value, +): + config = _make_config( + tmp_path, + hf_config=_official_config(**{field_name: value}), + ) + + assert getattr(config.hf_config, field_name) == value + + def test_minimax_config_requires_all_fp8_exclusions(tmp_path): hf_config = _official_config( quantization_config=_quantization_config( @@ -132,7 +154,7 @@ def test_minimax_config_rejects_unvalidated_parallel_layout( tmp_path, parallel_kwargs, ): - with pytest.raises(ValueError, match="MiniMax M2.7"): + with pytest.raises(ValueError, match="MiniMax M2.7|Outer-TP MoE"): _make_config(tmp_path, **parallel_kwargs) @@ -171,12 +193,23 @@ def _validate(method="", **overrides): "enable_prefix_caching": True, } values.update(overrides) + tp_size = values.pop("tensor_parallel_size") + ep_size = values.pop("expert_parallel_size") + dp_size = values.pop("data_parallel_size") + values["topology"] = ParallelTopology( + tp_size, + ep_size, + dp_size, + ParallelMode.OUTER_TP_MOE if tp_size > 1 else ParallelMode.STANDARD, + ) return validate_model_runtime_compatibility(**values) def test_minimax_compatibility_matches_qwen3_moe_sparse_runtime(): - assert MODEL_RUNTIME_COMPATIBILITY["minimax_m2"] is MINIMAX_M2_EP_COMPATIBILITY - assert MINIMAX_M2_EP_COMPATIBILITY.parallel_mode == "ep_replicated_kv" + assert ( + MODEL_RUNTIME_COMPATIBILITY["minimax_m2", ParallelMode.STANDARD] + is MINIMAX_M2_EP_COMPATIBILITY + ) assert MINIMAX_M2_EP_COMPATIBILITY.sparse_methods == { "", "streamingllm", @@ -200,7 +233,6 @@ def test_minimax_compatibility_matches_qwen3_moe_sparse_runtime(): def test_minimax_outer_tp_compatibility_preserves_sparse_matrix(): - assert MINIMAX_M2_TP_EP_COMPATIBILITY.parallel_mode == "outer_tp_moe_tp_ep" assert ( MINIMAX_M2_TP_EP_COMPATIBILITY.sparse_methods == MINIMAX_M2_EP_COMPATIBILITY.sparse_methods diff --git a/tests/test_moe_config.py b/tests/test_moe_config.py index 87e38c72..a60cf648 100644 --- a/tests/test_moe_config.py +++ b/tests/test_moe_config.py @@ -2,6 +2,7 @@ import torch from sparsevllm.triton_kernel.moe_config import ( + resolve_fp8_routed_gemm_config, resolve_moe_gemm_config, token_bucket, ) @@ -105,6 +106,96 @@ def test_h20_qwen3_moe_config_is_shape_and_stage_aware(): assert large.block_m == 64 +@pytest.mark.parametrize( + ("device_name", "num_local_experts", "intermediate_size", "expected"), + [ + ( + "NVIDIA H20", + 256, + 512, + { + "w13": ((32, 4, 4), (32, 4, 4), (64, 4, 3), (128, 8, 3)), + "gate_up_swiglu": ((32, 4, 4), (32, 4, 4), (32, 4, 3), (32, 4, 4)), + "w2": ((32, 4, 4), (64, 4, 3), (128, 4, 2), (64, 4, 4)), + }, + ), + ( + "NVIDIA H20", + 256, + 256, + { + "w13": ((32, 4, 4), (32, 4, 4), (32, 4, 4), (64, 4, 3)), + "gate_up_swiglu": ((32, 4, 4), (32, 4, 4), (32, 4, 4), (32, 4, 3)), + "w2": ((32, 4, 3), (64, 4, 2), (128, 4, 2), (128, 4, 2)), + }, + ), + ( + "NVIDIA H20", + 128, + 512, + { + "w13": ((32, 4, 4), (64, 4, 4), (64, 4, 4), (128, 8, 3)), + "gate_up_swiglu": ((32, 4, 4), (32, 4, 4), (32, 4, 4), (32, 4, 3)), + "w2": ((64, 4, 4), (32, 4, 3), (64, 4, 3), (64, 4, 4)), + }, + ), + ( + "NVIDIA H100 80GB HBM3", + 256, + 512, + { + "w13": ((32, 4, 4), (64, 4, 4), (64, 4, 3), (64, 4, 3)), + "gate_up_swiglu": ((32, 4, 4), (32, 4, 4), (32, 4, 3), (32, 4, 3)), + "w2": ((32, 4, 4), (32, 4, 4), (64, 4, 3), (32, 4, 3)), + }, + ), + ( + "NVIDIA H100 80GB HBM3", + 256, + 256, + { + "w13": ((32, 4, 4), (32, 4, 4), (32, 4, 4), (32, 4, 3)), + "gate_up_swiglu": ((32, 4, 4), (32, 4, 4), (32, 4, 4), (32, 4, 3)), + "w2": ((32, 4, 4), (64, 4, 4), (64, 4, 2), (32, 4, 3)), + }, + ), + ( + "NVIDIA H100 80GB HBM3", + 128, + 512, + { + "w13": ((32, 4, 4), (32, 4, 4), (32, 4, 4), (64, 4, 3)), + "gate_up_swiglu": ((32, 4, 4), (32, 4, 4), (32, 4, 4), (32, 4, 3)), + "w2": ((32, 4, 4), (32, 4, 4), (32, 4, 3), (32, 4, 3)), + }, + ), + ], +) +def test_qwen36_bf16_profiles_cover_decode_buckets( + device_name, num_local_experts, intermediate_size, expected +): + for stage, stage_expected in expected.items(): + configs = tuple( + resolve_moe_gemm_config( + dtype=torch.bfloat16, + num_tokens=tokens, + top_k=8, + num_local_experts=num_local_experts, + hidden_size=2048, + intermediate_size=intermediate_size, + stage=stage, + device_name=device_name, + device_capability=(9, 0), + ) + for tokens in (1, 2, 4, 8) + ) + assert tuple( + (config.block_n, config.num_warps, config.num_stages) + for config in configs + ) == stage_expected + assert all(config.block_m == 16 and config.block_k == 64 for config in configs) + + def test_fallback_heuristic_uses_logical_assignment_count(): common = dict( dtype=torch.float16, @@ -155,3 +246,106 @@ def test_h100_profile_switches_to_large_token_config(): ) assert resolve_moe_gemm_config(**common, num_tokens=512).block_m == 16 assert resolve_moe_gemm_config(**common, num_tokens=1024).block_m == 64 + + +@pytest.mark.parametrize( + ("device_name", "local_experts", "intermediate_size", "expected"), + [ + ( + "NVIDIA H20", + 256, + 512, + { + "w13": ((64, 4), (128, 4), (64, 4), (64, 3)), + "w2": ((64, 3), (64, 2), (64, 2), (128, 2)), + }, + ), + ( + "NVIDIA H20", + 256, + 256, + { + "w13": ((64, 4),) * 4, + "w2": ((64, 3), (64, 2), (64, 2), (128, 2)), + }, + ), + ( + "NVIDIA H20", + 128, + 512, + { + "w13": ((64, 4),) * 4, + "w2": ((64, 3), (128, 4), (64, 2), (64, 2)), + }, + ), + ( + "NVIDIA H100 80GB HBM3", + 256, + 512, + { + "w13": ((64, 5), (64, 4), (64, 3), (64, 4)), + "w2": ((64, 4), (128, 4), (64, 2), (64, 3)), + }, + ), + ( + "NVIDIA H100 80GB HBM3", + 256, + 256, + { + "w13": ((64, 4), (64, 5), (64, 4), (64, 3)), + "w2": ((64, 3), (64, 2), (64, 2), (64, 2)), + }, + ), + ( + "NVIDIA H100 80GB HBM3", + 128, + 512, + { + "w13": ((64, 4),) * 4, + "w2": ((64, 4), (64, 3), (64, 2), (64, 2)), + }, + ), + ], +) +def test_qwen36_fp8_profiles_cover_decode_buckets( + device_name, local_experts, intermediate_size, expected +): + for stage, stage_expected in expected.items(): + configs = tuple( + resolve_fp8_routed_gemm_config( + num_tokens=tokens, + top_k=8, + num_local_experts=local_experts, + hidden_size=2048, + intermediate_size=intermediate_size, + stage=stage, + device_name=device_name, + device_capability=(9, 0), + ) + for tokens in (1, 2, 4, 8) + ) + assert tuple( + (config.block_n, config.num_stages) for config in configs + ) == stage_expected + assert all( + config.block_m == 16 + and config.block_k == 128 + and config.num_warps == 4 + and config.swap_ab + for config in configs + ) + + +def test_fp8_routed_unknown_shape_uses_explicit_default(): + config = resolve_fp8_routed_gemm_config( + num_tokens=1, + top_k=8, + num_local_experts=64, + hidden_size=2048, + intermediate_size=512, + stage="w13", + device_name="NVIDIA H100 80GB HBM3", + device_capability=(9, 0), + ) + + assert (config.block_n, config.block_k, config.swap_ab) == (128, 128, False) diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 6e05b2e4..2805fd0c 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -5,6 +5,7 @@ import pytest import torch +from sparsevllm.operators.all_reduce import HopperTp2FlashInferAllReduceProvider from sparsevllm.operators.fp8_linear import ( FP8_LINEAR_REGISTRY, FlashInferSm90Fp8LinearProvider, @@ -12,7 +13,19 @@ TritonFp8LinearProvider, resolve_fp8_linear_provider, ) -from sparsevllm.operators.moe import MOE_REGISTRY, MoeOpSpec, resolve_moe_provider +from sparsevllm.operators.gate_up_swiglu import ( + GATE_UP_SWIGLU_REGISTRY, + GateUpSwiGLUOpSpec, + NativeGateUpSwiGLUProvider, +) +from sparsevllm.operators.moe import ( + MOE_REGISTRY, + FlashInferCutlassFp4MoeProvider, + FlashInferCutlassFp8MoeProvider, + HopperQwen36HybridFp8MoeProvider, + MoeOpSpec, + resolve_moe_provider, +) from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum @@ -64,6 +77,8 @@ def _moe_spec( tp_size=1, routing_method="softmax", scale_dtype=None, + cuda_graph=True, + activation_limit=None, ) -> MoeOpSpec: return MoeOpSpec( num_experts=num_experts, @@ -75,10 +90,11 @@ def _moe_spec( weight_dtype=weight_dtype, block_shape=block_shape, ep_size=ep_size, - cuda_graph=True, + cuda_graph=cuda_graph, tp_size=tp_size, routing_method=routing_method, scale_dtype=scale_dtype, + activation_limit=activation_limit, ) @@ -97,6 +113,88 @@ def _linear_spec( ) +def _gate_up_spec(**overrides) -> GateUpSwiGLUOpSpec: + values = { + "hidden_size": 2048, + "intermediate_size": 512, + "tp_size": 1, + "activation_dtype": torch.bfloat16, + "weight_dtype": torch.bfloat16, + "cuda_graph": True, + } + values.update(overrides) + return GateUpSwiGLUOpSpec(**values) + + +def test_flashinfer_all_reduce_falls_back_before_unsupported_shape_launch(): + provider = HopperTp2FlashInferAllReduceProvider.__new__( + HopperTp2FlashInferAllReduceProvider + ) + provider.fallback = Mock() + tensor = torch.randn(1, 248320, dtype=torch.bfloat16) + provider.fallback.run.return_value = tensor + + assert provider.run(tensor) is tensor + provider.fallback.run.assert_called_once_with(tensor) + + +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_h20_gate_up_provider_accepts_profiled_qwen36_shape(tp_size): + resolved = OpResolver(GATE_UP_SWIGLU_REGISTRY).resolve( + _gate_up_spec(tp_size=tp_size), + _cuda_caps((9, 0), device_name="NVIDIA H20"), + ) + + assert resolved.provider.name == "h20_triton_decode" + + +@pytest.mark.parametrize( + "spec,caps", + [ + ( + _gate_up_spec(weight_dtype=torch.float8_e4m3fn), + _cuda_caps((9, 0), device_name="NVIDIA H20"), + ), + ( + _gate_up_spec(intermediate_size=768), + _cuda_caps((9, 0), device_name="NVIDIA H20"), + ), + ( + _gate_up_spec(), + _cuda_caps((9, 0), device_name="NVIDIA H100 80GB HBM3"), + ), + (_gate_up_spec(), _cuda_caps((8, 9), device_name="NVIDIA H20")), + ], +) +def test_unprofiled_gate_up_shape_uses_native_provider(spec, caps): + resolved = OpResolver(GATE_UP_SWIGLU_REGISTRY).resolve(spec, caps) + + assert resolved.provider.name == "native" + + +def test_native_gate_up_provider_matches_swiglu_semantics(): + torch.manual_seed(0) + inputs = torch.randn(3, 8) + projection = torch.nn.Linear(8, 12, bias=False) + with torch.inference_mode(): + projected = projection(inputs) + gate, up = projected.chunk(2, dim=-1) + expected = torch.nn.functional.silu(gate) * up + actual = NativeGateUpSwiGLUProvider().run( + _gate_up_spec( + hidden_size=8, + intermediate_size=6, + activation_dtype=torch.float32, + weight_dtype=torch.float32, + cuda_graph=False, + ), + inputs, + projection, + ) + + torch.testing.assert_close(actual, expected) + + @pytest.mark.parametrize( "overrides", [ @@ -155,9 +253,9 @@ def test_fp8_linear_uses_generic_triton_when_specialization_does_not_match( def test_fp8_linear_prefers_flashinfer_on_sm90(): - with ( - patch("sparsevllm.operators.fp8_linear.find_spec", return_value=object()), - patch("sparsevllm.operators.fp8_linear.which", return_value="/cuda/bin/nvcc"), + with patch( + "sparsevllm.operators.fp8_linear.find_spec", + return_value=object(), ): resolved = OpResolver(FP8_LINEAR_REGISTRY).resolve( _linear_spec(), @@ -179,9 +277,9 @@ def test_fp8_linear_uses_triton_when_flashinfer_is_missing_on_sm90(): def test_fp8_linear_resolution_does_not_require_nvcc(): - with ( - patch("sparsevllm.operators.fp8_linear.find_spec", return_value=object()), - patch("sparsevllm.operators.fp8_linear.which", return_value=None), + with patch( + "sparsevllm.operators.fp8_linear.find_spec", + return_value=object(), ): resolved = OpResolver(FP8_LINEAR_REGISTRY).resolve( _linear_spec(), @@ -191,13 +289,12 @@ def test_fp8_linear_resolution_does_not_require_nvcc(): assert resolved.provider.name == "flashinfer_sm90" -def test_flashinfer_linear_binds_triton_for_missing_uncached_kernel(): +def test_flashinfer_linear_does_not_mask_missing_jit_artifact(): flashinfer_call = Mock( side_effect=RuntimeError( "Assertion failed: !cubin.empty() || isPathValid(path_)" ) ) - fallback_output = torch.ones(2, 128, dtype=torch.bfloat16) provider = FlashInferSm90Fp8LinearProvider() x = torch.ones(2, 128, dtype=torch.bfloat16) weight = torch.ones(128, 128).to(torch.float8_e4m3fn) @@ -212,20 +309,11 @@ def test_flashinfer_linear_binds_triton_for_missing_uncached_kernel(): ) }, ), - patch("sparsevllm.operators.fp8_linear.which", return_value=None), - patch.object( - TritonFp8LinearProvider, - "__call__", - return_value=fallback_output, - ) as triton_call, + pytest.raises(RuntimeError, match="cubin.empty"), ): - first = provider(x, weight, scale) - second = provider(x, weight, scale) + provider(x, weight, scale) - assert first is fallback_output - assert second is fallback_output assert flashinfer_call.call_count == 1 - assert triton_call.call_count == 2 def test_flashinfer_linear_does_not_mask_other_runtime_failures(): @@ -244,14 +332,10 @@ def test_flashinfer_linear_does_not_mask_other_runtime_failures(): ) }, ), - patch("sparsevllm.operators.fp8_linear.which", return_value=None), pytest.raises(RuntimeError, match="invalid scale layout"), ): provider(x, weight, scale) - assert provider._fallback is None - - def test_fp8_linear_reports_unsupported_pre_fp8_device(): with pytest.raises(RuntimeError, match="native FP8 tensor cores"): OpResolver(FP8_LINEAR_REGISTRY).resolve( @@ -372,6 +456,44 @@ def test_hopper_fused_moe_uses_profiled_tp_ep_shape(): assert resolved.provider.name == "triton_hopper_fused" +@pytest.mark.parametrize( + ("tp_size", "ep_size", "intermediate_size", "num_local_experts"), + [(1, 1, 512, 256), (2, 1, 256, 256), (1, 2, 512, 128)], +) +def test_qwen36_bf16_moe_uses_profiled_hopper_provider( + tp_size, + ep_size, + intermediate_size, + num_local_experts, +): + spec = _moe_spec( + activation_dtype=torch.bfloat16, + weight_dtype=torch.bfloat16, + block_shape=None, + hidden_size=2048, + intermediate_size=intermediate_size, + num_local_experts=num_local_experts, + num_experts=256, + top_k=8, + ep_size=ep_size, + tp_size=tp_size, + ) + + resolved = OpResolver(MOE_REGISTRY).resolve( + spec, + _cuda_caps((9, 0), native_fp8=False, device_name="NVIDIA H20"), + ) + + assert resolved.provider.name == "h20_qwen36_fused_bf16" + h100 = OpResolver(MOE_REGISTRY).resolve( + spec, + _cuda_caps( + (9, 0), native_fp8=False, device_name="NVIDIA H100 80GB HBM3" + ), + ) + assert h100.provider.name == "triton_hopper_fused" + + @pytest.mark.parametrize( ("tp_size", "ep_size", "intermediate_size", "num_local_experts"), [(4, 1, 384, 256), (2, 2, 768, 128), (1, 4, 1536, 64)], @@ -520,6 +642,317 @@ def test_fp8_moe_prefers_flashinfer_only_on_sm90(): assert blackwell.provider.name == "triton" +def test_fp4_moe_prefers_flashinfer_on_sm90(): + spec = _moe_spec( + hidden_size=4096, + intermediate_size=2048, + num_experts=256, + num_local_experts=64, + top_k=6, + ep_size=4, + weight_dtype=torch.uint8, + scale_dtype=torch.uint8, + block_shape=(1, 32), + activation_limit=10.0, + ) + with patch("sparsevllm.operators.moe.find_spec", return_value=object()): + resolved = OpResolver(MOE_REGISTRY).resolve(spec, _cuda_caps((9, 0))) + + assert resolved.provider.name == "flashinfer_cutlass_fp4_sm90" + + +def test_fp4_moe_loader_preserves_packed_bytes(): + provider = FlashInferCutlassFp4MoeProvider() + spec = _moe_spec( + hidden_size=256, + intermediate_size=128, + weight_dtype=torch.uint8, + scale_dtype=torch.uint8, + block_shape=(1, 32), + activation_limit=10.0, + ) + w13 = torch.zeros(4, 256, 128, dtype=torch.uint8) + w2 = torch.zeros(4, 256, 64, dtype=torch.uint8) + s13 = torch.zeros(4, 256, 8, dtype=torch.uint8) + s2 = torch.zeros(4, 256, 4, dtype=torch.uint8) + weight_bits = torch.arange(128 * 128, dtype=torch.int64).to(torch.uint8) + weight_bits = weight_bits.view(128, 128) + scale_bits = torch.arange(128 * 8, dtype=torch.int64).to(torch.uint8) + scale_bits = scale_bits.view(128, 8) + + provider.load_expert_projection( + spec, + local_expert_id=1, + projection="up", + loaded_weight=weight_bits.view(torch.int8), + loaded_scale=scale_bits.view(torch.float8_e8m0fnu), + w13_weight=w13, + w2_weight=w2, + w13_scale_inv=s13, + w2_scale_inv=s2, + ) + + assert torch.equal(w13[1, :128], weight_bits) + assert torch.equal(s13[1, :128], scale_bits) + + +def test_qwen36_hybrid_moe_uses_profiled_graph_shape_on_h100(): + spec = _moe_spec( + hidden_size=2048, + intermediate_size=512, + num_local_experts=128, + num_experts=256, + top_k=8, + ep_size=2, + tp_size=1, + ) + with patch("sparsevllm.operators.moe.find_spec", return_value=object()): + resolved = OpResolver(MOE_REGISTRY).resolve( + spec, + _cuda_caps((9, 0), device_name="NVIDIA H100 80GB HBM3"), + ) + + assert resolved.provider.name == "hopper_qwen36_hybrid_fp8" + assert resolved.provider.gate_up_order == "up_gate" + + +def test_qwen36_hybrid_moe_uses_profiled_single_gpu_shape_on_h100(): + spec = _moe_spec( + hidden_size=2048, + intermediate_size=512, + num_local_experts=256, + num_experts=256, + top_k=8, + ep_size=1, + tp_size=1, + ) + with patch("sparsevllm.operators.moe.find_spec", return_value=object()): + resolved = OpResolver(MOE_REGISTRY).resolve( + spec, + _cuda_caps((9, 0), device_name="NVIDIA H100 80GB HBM3"), + ) + + assert resolved.provider.name == "hopper_qwen36_hybrid_fp8" + assert resolved.provider.gate_up_order == "up_gate" + + +def test_qwen36_pure_tp_uses_triton_for_sharded_experts(): + spec = _moe_spec( + hidden_size=2048, + intermediate_size=256, + num_local_experts=256, + num_experts=256, + top_k=8, + ep_size=1, + tp_size=2, + ) + with patch("sparsevllm.operators.moe.find_spec", return_value=object()): + resolved = OpResolver(MOE_REGISTRY).resolve( + spec, + _cuda_caps((9, 0), device_name="NVIDIA H100 80GB HBM3"), + ) + + assert resolved.provider.name == "triton" + + +@pytest.mark.parametrize( + ("spec_overrides", "caps_overrides", "reason"), + [ + ( + {}, + {"supports_graph_capture": False}, + "device does not support CUDA Graph capture", + ), + ({"intermediate_size": 640}, {}, "requires profiled Qwen3.6"), + ], +) +def test_qwen36_hybrid_moe_rejects_unprofiled_execution( + spec_overrides, + caps_overrides, + reason, +): + values = dict( + hidden_size=2048, + intermediate_size=512, + num_local_experts=128, + num_experts=256, + top_k=8, + ep_size=2, + tp_size=1, + ) + values.update(spec_overrides) + spec = _moe_spec(**values) + caps = _cuda_caps((9, 0), device_name="NVIDIA H100 80GB HBM3") + caps = DeviceCaps(**{**caps.__dict__, **caps_overrides}) + + with patch("sparsevllm.operators.moe.find_spec", return_value=object()): + resolved = OpResolver(MOE_REGISTRY).resolve(spec, caps) + + assert resolved.provider.name == "flashinfer_cutlass_fp8_sm90" + assert reason in dict(resolved.rejected)["hopper_qwen36_hybrid_fp8"] + + +def test_qwen36_hybrid_moe_supports_eager_execution(): + spec = _moe_spec( + hidden_size=2048, + intermediate_size=512, + num_local_experts=256, + num_experts=256, + top_k=8, + ep_size=1, + tp_size=1, + cuda_graph=False, + ) + with patch("sparsevllm.operators.moe.find_spec", return_value=object()): + resolved = OpResolver(MOE_REGISTRY).resolve( + spec, + _cuda_caps((9, 0), device_name="NVIDIA H100 80GB HBM3"), + ) + + assert resolved.provider.name == "hopper_qwen36_hybrid_fp8" + + +def test_h20_qwen36_hybrid_moe_uses_profiled_provider(): + spec = _moe_spec( + hidden_size=2048, + intermediate_size=512, + num_local_experts=256, + num_experts=256, + top_k=8, + ep_size=1, + tp_size=1, + ) + with patch("sparsevllm.operators.moe.find_spec", return_value=object()): + resolved = OpResolver(MOE_REGISTRY).resolve( + spec, + _cuda_caps((9, 0), device_name="NVIDIA H20"), + ) + + assert resolved.provider.name == "h20_qwen36_hybrid_fp8" + + +def test_qwen36_hybrid_moe_dispatches_by_token_bucket(): + provider = HopperQwen36HybridFp8MoeProvider() + spec = _moe_spec( + hidden_size=2048, + intermediate_size=512, + num_local_experts=128, + num_experts=256, + top_k=8, + ep_size=2, + tp_size=1, + ) + small_output = torch.ones(4, 2) + large_output = torch.ones(5, 2) + triton_call = Mock(return_value=small_output) + weights = torch.empty(1) + + with patch.dict( + sys.modules, + { + "sparsevllm.triton_kernel.moe": SimpleNamespace( + fused_moe_fp8=triton_call + ) + }, + ): + actual_small = provider.run( + spec, + torch.empty(4, 2), + torch.empty(4, 8, dtype=torch.int32), + torch.empty(4, 8), + weights, + weights, + weights, + weights, + local_expert_start=128, + ep_rank=1, + ) + + with patch.object( + FlashInferCutlassFp8MoeProvider, + "run", + return_value=large_output, + ) as flashinfer_call: + actual_large = provider.run( + spec, + torch.empty(5, 2), + torch.empty(5, 8, dtype=torch.int32), + torch.empty(5, 8), + weights, + weights, + weights, + weights, + local_expert_start=128, + ep_rank=1, + ) + + assert actual_small is small_output + assert actual_large is large_output + assert triton_call.call_args.kwargs["gate_up_order"] == "up_gate" + flashinfer_call.assert_called_once() + + +def test_qwen36_hybrid_moe_uses_larger_triton_bucket_on_single_gpu(): + provider = HopperQwen36HybridFp8MoeProvider() + spec = _moe_spec( + hidden_size=2048, + intermediate_size=512, + num_local_experts=256, + num_experts=256, + top_k=8, + ep_size=1, + tp_size=1, + ) + triton_output = torch.ones(8, 2) + flashinfer_output = torch.ones(9, 2) + triton_call = Mock(return_value=triton_output) + weights = torch.empty(1) + + with patch.dict( + sys.modules, + { + "sparsevllm.triton_kernel.moe": SimpleNamespace( + fused_moe_fp8=triton_call + ) + }, + ): + actual_decode = provider.run( + spec, + torch.empty(8, 2), + torch.empty(8, 8, dtype=torch.int32), + torch.empty(8, 8), + weights, + weights, + weights, + weights, + local_expert_start=0, + ep_rank=0, + ) + + with patch.object( + FlashInferCutlassFp8MoeProvider, + "run", + return_value=flashinfer_output, + ) as flashinfer_call: + actual_prefill = provider.run( + spec, + torch.empty(9, 2), + torch.empty(9, 8, dtype=torch.int32), + torch.empty(9, 8), + weights, + weights, + weights, + weights, + local_expert_start=0, + ep_rank=0, + ) + + assert actual_decode is triton_output + assert actual_prefill is flashinfer_output + assert triton_call.call_args.kwargs["gate_up_order"] == "up_gate" + flashinfer_call.assert_called_once() + + def test_fp8_moe_uses_triton_when_flashinfer_is_missing_on_sm90(): with patch("sparsevllm.operators.moe.find_spec", return_value=None): resolved = OpResolver(MOE_REGISTRY).resolve( diff --git a/tests/test_operator_registry.py b/tests/test_operator_registry.py index d3cd4573..2070f828 100644 --- a/tests/test_operator_registry.py +++ b/tests/test_operator_registry.py @@ -1,8 +1,16 @@ from dataclasses import dataclass +from unittest.mock import patch import pytest -from sparsevllm.operators.registry import OpRegistry, OpResolver, SupportResult +import sparsevllm.operators.registry as operator_registry +from sparsevllm.operators.registry import ( + OpRegistry, + OpResolver, + SupportResult, + log_operator_implementations, + record_operator_binding, +) from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum @@ -42,7 +50,9 @@ class Specialized: def supports(cls, spec, caps): return SupportResult.yes() if spec.enabled else SupportResult.no("disabled") - resolved = OpResolver(registry).resolve(_Spec(), _caps()) + with patch.dict(operator_registry._OPERATOR_BINDINGS, {}, clear=True): + resolved = OpResolver(registry).resolve(_Spec(), _caps()) + assert resolved.provider in operator_registry._OPERATOR_BINDINGS["_test"] assert resolved.provider.name == "specialized" @@ -149,6 +159,32 @@ def supports(cls, spec, caps): assert resolved.rejected == (("specialized", "missing optional library"),) +def test_operator_organization_logs_live_bound_implementations(): + class Provider: + def __init__(self, implementation_name): + self.implementation_name = implementation_name + + attention = Provider("triton") + linear = Provider("flashinfer_sm90") + linear_fallback = Provider("triton") + + with ( + patch.dict(operator_registry._OPERATOR_BINDINGS, {}, clear=True), + patch("sparsevllm.operators.registry.logger.info") as log_info, + ): + record_operator_binding("Attention", attention) + record_operator_binding("block-scaled FP8 Linear", linear) + record_operator_binding("block-scaled FP8 Linear", linear_fallback) + + log_operator_implementations() + + log_info.assert_called_once_with( + "Operator implementations:\n{}", + " Attention: triton\n" + " block-scaled FP8 Linear: flashinfer_sm90, triton", + ) + + def test_resolver_forwards_provider_constructor_arguments(): registry = OpRegistry("_test") diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py index 76bfa1bc..e2d91b2c 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -1,16 +1,18 @@ from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch import torch.distributed as dist import sparsevllm.platforms as platforms -from sparsevllm.config import Config -from sparsevllm.distributed import ParallelContext, ParallelGroup -from sparsevllm.distributed.parallel_context import ( +from sparsevllm.config import Config, RuntimeLayout +from sparsevllm.distributed import ( + ParallelContext, + ParallelGroup, + ParallelMode, + ParallelTopology, get_parallel_context, - hybrid_moe_group_ranks, init_parallel_context, parallel_group_ranks, parallel_ranks_from_world_rank, @@ -88,51 +90,78 @@ def _hf_config(model_type: str = "qwen3_moe", *, num_experts: int = 8): def test_world_rank_mapping_round_trips(): + topology = ParallelTopology(2, 3, 4) for world_rank in range(24): - ranks = parallel_ranks_from_world_rank( - world_rank, - tp_size=2, - ep_size=3, - dp_size=4, - ) - assert world_rank_from_parallel_ranks( - *ranks, - tp_size=2, - ep_size=3, - dp_size=4, - ) == world_rank + ranks = parallel_ranks_from_world_rank(topology, world_rank) + assert world_rank_from_parallel_ranks(topology, *ranks) == world_rank def test_parallel_group_members_follow_dp_ep_tp_layout(): - assert parallel_group_ranks(tp_size=2, ep_size=2, dp_size=2) == { - "tensor": ((0, 1), (2, 3), (4, 5), (6, 7)), + tensor_groups = ((0, 1), (2, 3), (4, 5), (6, 7)) + assert parallel_group_ranks(ParallelTopology(2, 2, 2)) == { + "tensor": tensor_groups, "expert": ((0, 2), (1, 3), (4, 6), (5, 7)), "data": ((0, 4), (1, 5), (2, 6), (3, 7)), + "moe_tensor": tensor_groups, } def test_hybrid_moe_groups_split_outer_attention_world(): - assert hybrid_moe_group_ranks(outer_tp_size=4, moe_ep_size=2) == { - "attention": ((0, 1, 2, 3),), + topology = ParallelTopology(4, 2, 1, ParallelMode.OUTER_TP_MOE) + assert parallel_group_ranks(topology) == { + "tensor": ((0, 1, 2, 3),), "moe_tensor": ((0, 1), (2, 3)), - "moe_expert": ((0, 2), (1, 3)), + "expert": ((0, 2), (1, 3)), "data": ((0,), (1,), (2,), (3,)), } +def test_dpa_ep_groups_overlap_data_and_expert_worlds(): + topology = ParallelTopology(1, 4, 4, ParallelMode.DPA_EP) + assert parallel_group_ranks(topology) == { + "tensor": ((0,), (1,), (2,), (3,)), + "moe_tensor": ((0,), (1,), (2,), (3,)), + "expert": ((0, 1, 2, 3),), + "data": ((0, 1, 2, 3),), + } + + +def test_parallel_topology_resolves_rank_local_sizes(): + standard = ParallelTopology(2, 4, 1) + hybrid = ParallelTopology(4, 2, 1, ParallelMode.OUTER_TP_MOE) + dpa = ParallelTopology(1, 4, 4, ParallelMode.DPA_EP) + + assert (standard.world_size, standard.attention_tp_size, standard.moe_tp_size) == (8, 2, 2) + assert (hybrid.world_size, hybrid.attention_tp_size, hybrid.moe_tp_size) == (4, 4, 2) + assert (dpa.world_size, dpa.attention_tp_size, dpa.moe_tp_size) == (4, 1, 1) + + +@pytest.mark.parametrize( + "topology", + [ + (0, 1, 1, ParallelMode.STANDARD), + (4, 3, 1, ParallelMode.OUTER_TP_MOE), + (4, 2, 2, ParallelMode.OUTER_TP_MOE), + (2, 2, 2, ParallelMode.DPA_EP), + (1, 2, 3, ParallelMode.DPA_EP), + ], +) +def test_parallel_topology_rejects_invalid_sizes(topology): + with pytest.raises(ValueError): + ParallelTopology(*topology) + + def test_hybrid_moe_parallel_context_uses_explicit_groups(): reset_parallel_context() with ( patch.object(dist, "is_initialized", return_value=True), patch.object(dist, "get_world_size", return_value=4), patch.object(dist, "get_rank", return_value=2), + patch.object(dist, "get_backend", return_value=dist.Backend.GLOO), patch.object(dist, "new_group", side_effect=lambda _ranks: object()), ): context = init_parallel_context( - tp_size=4, - ep_size=2, - dp_size=1, - hybrid_moe=True, + topology=ParallelTopology(4, 2, 1, ParallelMode.OUTER_TP_MOE), ) assert context.attention.ranks == (0, 1, 2, 3) assert context.attention_tp_rank == 2 @@ -143,6 +172,27 @@ def test_hybrid_moe_parallel_context_uses_explicit_groups(): reset_parallel_context() +def test_dpa_ep_parallel_context_uses_overlapping_groups(): + reset_parallel_context() + with ( + patch.object(dist, "is_initialized", return_value=True), + patch.object(dist, "get_world_size", return_value=4), + patch.object(dist, "get_rank", return_value=2), + patch.object(dist, "get_backend", return_value=dist.Backend.GLOO), + patch.object(dist, "new_group", side_effect=lambda _ranks: object()), + ): + context = init_parallel_context( + topology=ParallelTopology(1, 4, 4, ParallelMode.DPA_EP), + ) + assert context.attention.ranks == (2,) + assert context.expert.ranks == (0, 1, 2, 3) + assert context.data.ranks == (0, 1, 2, 3) + assert context.moe_tensor.ranks == (2,) + assert context.tp_rank == context.moe_tp_rank == 0 + assert context.dp_rank == context.ep_rank == 2 + reset_parallel_context() + + def test_parallel_context_lifecycle_and_local_groups(): reset_parallel_context() fake_groups = [] @@ -156,9 +206,11 @@ def new_group(ranks): patch.object(dist, "is_initialized", return_value=True), patch.object(dist, "get_world_size", return_value=4), patch.object(dist, "get_rank", return_value=2), + patch.object(dist, "get_backend", return_value=dist.Backend.GLOO), patch.object(dist, "new_group", side_effect=new_group), ): - context = init_parallel_context(tp_size=1, ep_size=2, dp_size=2) + topology = ParallelTopology(1, 2, 2) + context = init_parallel_context(topology=topology) assert context.world_rank == 2 assert context.tp_rank == 0 assert context.tp_size == 1 @@ -168,7 +220,7 @@ def new_group(ranks): assert context.data.ranks == (0, 2) assert get_parallel_context() is context with pytest.raises(RuntimeError, match="already initialized"): - init_parallel_context(tp_size=1, ep_size=2, dp_size=2) + init_parallel_context(topology=topology) assert [ranks for ranks, _ in fake_groups] == [ (0, 1), @@ -189,7 +241,7 @@ def test_parallel_context_rejects_world_size_mismatch(): patch.object(dist, "get_rank", return_value=0), ): with pytest.raises(ValueError, match="does not match"): - init_parallel_context(tp_size=1, ep_size=4, dp_size=1) + init_parallel_context(topology=ParallelTopology(1, 4, 1)) def test_ep_broadcast_uses_source_world_rank(): @@ -235,7 +287,7 @@ def test_qwen3_moe_parallel_config_validation(tmp_path): assert config.expert_parallel_size == 1 with patch("sparsevllm.configs.runtime.AutoConfig.from_pretrained", return_value=_hf_config()): - with pytest.raises(ValueError, match="num_key_value_heads must be divisible"): + with pytest.raises(ValueError, match="num_key_value_heads"): Config(model=str(tmp_path), tensor_parallel_size=4) fp16 = _hf_config() @@ -245,7 +297,7 @@ def test_qwen3_moe_parallel_config_validation(tmp_path): Config(model=str(tmp_path), tensor_parallel_size=2) with patch("sparsevllm.configs.runtime.AutoConfig.from_pretrained", return_value=_hf_config()): - with pytest.raises(ValueError, match="outer tensor_parallel_size"): + with pytest.raises(ValueError, match="TP divisible by EP"): Config(model=str(tmp_path), tensor_parallel_size=3, expert_parallel_size=2) with patch("sparsevllm.configs.runtime.AutoConfig.from_pretrained", return_value=_hf_config(num_experts=6)): @@ -352,7 +404,7 @@ def test_qwen3_dense_fp8_rejects_wrong_architecture(tmp_path): def test_dense_config_rejects_expert_or_data_parallelism(tmp_path): with patch("sparsevllm.configs.runtime.AutoConfig.from_pretrained", return_value=_hf_config("qwen3")): - with pytest.raises(ValueError, match="requires EP=1 and DP=1"): + with pytest.raises(ValueError, match="does not support expert parallelism"): Config(model=str(tmp_path), expert_parallel_size=2) @@ -371,6 +423,24 @@ def test_dense_layers_use_tp_group_in_replicated_ep_topology(): assert embedding.weight.shape == (32, 8) +def test_parallel_group_uses_bound_all_reduce_provider(): + provider = Mock() + input_tensor, output_tensor = torch.randn(2), torch.randn(2) + provider.run.return_value = output_tensor + group = ParallelGroup( + process_group=None, + ranks=(0, 1), + rank=0, + size=2, + all_reduce_provider=provider, + ) + + actual = ParallelContext._all_reduce(input_tensor, group) + + assert actual is output_tensor + provider.run.assert_called_once_with(input_tensor) + + def test_cache_kv_heads_depend_on_tp_not_ep(): context = _replicated_ep_context() config = SimpleNamespace( @@ -381,7 +451,7 @@ def test_cache_kv_heads_depend_on_tp_not_ep(): hidden_size=32, head_dim=4, ), - runtime_layout=None, + runtime_layout=RuntimeLayout.dense(2), max_model_len=128, max_num_seqs_in_gpu=2, max_num_seqs_in_batch=2, diff --git a/tests/test_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index f1f44396..969359a6 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -831,7 +831,7 @@ def test_h2o_accepts_supported_dense_model_tp_and_rejects_unknown_model(self): ): with self.assertRaisesRegex( NotImplementedError, - "model types already implemented", + "Unsupported Sparse-vLLM model_type", ): Config(model=str(Path(tmp)), vllm_sparse_method="h2o") @@ -1595,6 +1595,7 @@ def test_graph_warmup_uses_distinct_prompts_across_requests_and_rounds(self): max_decoding_seqs=3, max_model_len=2048, hf_config=SimpleNamespace(vocab_size=32), + model_spec=SimpleNamespace(num_experts_field=None), ) prompts = [] pending = 0 @@ -1639,6 +1640,7 @@ def step(): [ ("set_warmup_fake_prefill_attention", True, 2046), ("set_warmup_fake_prefill_attention", False), + ("log_operator_implementations",), ], ) @@ -1685,7 +1687,7 @@ def test_moe_workspace_warmup_profiles_decode_and_maximum_mlp_shapes(self): config.max_decoding_seqs = 24 config.max_num_batched_tokens = 56_214 config.mlp_chunk_size = 16_384 - config.hf_config = SimpleNamespace(model_type="qwen3_moe") + config.model_spec = SimpleNamespace(num_experts_field="num_experts") self.assertEqual( _moe_workspace_warmup_token_counts(config), @@ -1694,7 +1696,7 @@ def test_moe_workspace_warmup_profiles_decode_and_maximum_mlp_shapes(self): def test_dense_model_skips_moe_workspace_warmup(self): config = self.make_config(method="vanilla") - config.hf_config = SimpleNamespace(model_type="qwen2") + config.model_spec = SimpleNamespace(num_experts_field=None) self.assertEqual(_moe_workspace_warmup_token_counts(config), ()) @@ -1704,7 +1706,7 @@ def test_engine_runs_each_moe_workspace_shape_after_regular_warmup(self): max_decoding_seqs=24, max_num_batched_tokens=56_214, mlp_chunk_size=16_384, - hf_config=SimpleNamespace(model_type="qwen3_moe"), + model_spec=SimpleNamespace(num_experts_field="num_experts"), ) calls = [] engine.model_runner = SimpleNamespace( @@ -1727,7 +1729,7 @@ def test_moe_workspace_oom_fails_startup(self): max_decoding_seqs=24, max_num_batched_tokens=56_214, mlp_chunk_size=16_384, - hf_config=SimpleNamespace(model_type="qwen3_moe"), + model_spec=SimpleNamespace(num_experts_field="num_experts"), ) def fail_on_workspace(_method, _num_tokens): diff --git a/tests/test_qwen35_mixed_runtime.py b/tests/test_qwen35_mixed_runtime.py index 24ae8118..75a54a29 100644 --- a/tests/test_qwen35_mixed_runtime.py +++ b/tests/test_qwen35_mixed_runtime.py @@ -1,15 +1,15 @@ import json from collections import deque from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch import sparsevllm.platforms as platforms from sparsevllm.platforms import device_runtime -from sparsevllm.config import Config, RuntimeLayout -from sparsevllm.distributed import ParallelContext, ParallelGroup +from sparsevllm.config import Config, QuantizationConfig, RuntimeLayout +from sparsevllm.distributed import ParallelContext, ParallelGroup, ParallelTopology from sparsevllm.engine.cache_manager.base import ( CacheManager, LayerBatchStates, @@ -40,10 +40,14 @@ from sparsevllm.engine.sparse_controller import LayerBatchSparseState, SparseController from sparsevllm.models.qwen3_5 import ( Qwen35ForCausalLM, + Qwen35LinearAttention, Qwen35LinearConv1D, Qwen35RMSNorm, _get_rotary_dim, ) +from sparsevllm.models.qwen3_5_moe import Qwen35MoeRouter, Qwen35MoeSparseMoeBlock +from sparsevllm.models.checkpoint import validate_checkpoint +from sparsevllm.models.spec import resolve_model_spec from sparsevllm.platforms.cpu import CpuPlatform from sparsevllm.sampling_params import SamplingParams from sparsevllm.utils.loader import _target_weight_name_for_model, _validate_all_quantized_weights_loaded @@ -96,6 +100,185 @@ def _make_config(tmp_path, **kwargs): return Config(model=str(tmp_path), **kwargs) +def test_qwen35_moe_checkpoint_validation_accepts_architecture_variants(): + hf_config = SimpleNamespace( + vocab_size=32000, + hidden_size=1536, + num_hidden_layers=8, + num_experts=128, + num_experts_per_tok=4, + layer_types=["linear_attention", "full_attention"] * 4, + torch_dtype=torch.bfloat16, + hidden_act="silu", + attn_output_gate=True, + attention_bias=False, + partial_rotary_factor=0.25, + mamba_ssm_dtype="float32", + rms_norm_eps=1.0e-6, + tie_word_embeddings=False, + ) + + validate_checkpoint( + "qwen3_5_moe", + outer_config=SimpleNamespace( + architectures=["Qwen3_5MoeForConditionalGeneration"] + ), + config=hf_config, + raw_quantization_config=None, + quantization=QuantizationConfig.disabled(model_name="Qwen3.6 MoE"), + topology=ParallelTopology(1, 1, 1), + ) + + +def test_linear_attention_fuses_qkvz_and_ba_projections(): + config = _qwen35_outer_config(num_layers=1, full_layers=()).text_config + config.quantization_config = None + config.runtime_recurrent_state_dtype = torch.float32 + parallel_context = _single_process_parallel_context() + with ( + patch( + "sparsevllm.models.qwen3_5.get_parallel_context", + return_value=parallel_context, + ), + patch( + "sparsevllm.layers.linear.get_parallel_context", + return_value=parallel_context, + ), + ): + attention = Qwen35LinearAttention(config) + + torch.manual_seed(0) + attention.in_proj_qkvz.weight.data.normal_() + attention.in_proj_ba.weight.data.normal_() + hidden_states = torch.randn(3, config.hidden_size) + qkvz = torch.nn.functional.linear( + hidden_states, + attention.in_proj_qkvz.weight, + ) + q, k, v, z = qkvz.split( + [ + attention.tp_key_dim, + attention.tp_key_dim, + attention.tp_value_dim, + attention.tp_value_dim, + ], + dim=-1, + ) + expected_b, expected_a = torch.nn.functional.linear( + hidden_states, + attention.in_proj_ba.weight, + ).chunk(2, dim=-1) + + mixed_qkv, actual_z, actual_b, actual_a = attention._project_qkvzba( + hidden_states + ) + + torch.testing.assert_close(mixed_qkv, torch.cat((q, k, v), dim=-1)) + torch.testing.assert_close( + actual_z, + z.view(-1, attention.num_v_heads, attention.head_v_dim), + ) + torch.testing.assert_close(actual_b, expected_b) + torch.testing.assert_close(actual_a, expected_a) + assert Qwen35ForCausalLM.packed_modules_mapping["in_proj_z"] == ( + "in_proj_qkvz", + 3, + ) + + +def test_qwen35_moe_fuses_router_and_shared_gate_projection(): + router = Qwen35MoeRouter.__new__(Qwen35MoeRouter) + torch.nn.Module.__init__(router) + router.num_experts, router.hidden_size = 4, 3 + router.weight = torch.nn.Parameter(torch.empty(5, 3)) + router.op_spec = object() + topk_weights, topk_ids = torch.ones(2, 1), torch.zeros(2, 1, dtype=torch.int32) + router.provider = Mock() + router.provider.run.return_value = (topk_weights, topk_ids) + expert_weight = torch.arange(12, dtype=torch.float32).view(4, 3) + shared_weight = torch.tensor([[2.0, 3.0, 5.0]]) + router.weight_loader(router.weight, expert_weight) + router.weight_loader(router.weight, shared_weight, "shared") + hidden_states = torch.tensor([[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]]) + + actual_weights, actual_ids, shared_logits = router(hidden_states) + + assert actual_weights is topk_weights + assert actual_ids is topk_ids + torch.testing.assert_close( + shared_logits, + torch.nn.functional.linear(hidden_states, shared_weight), + ) + router.provider.run.assert_called_once() + assert Qwen35ForCausalLM.packed_modules_mapping["shared_expert_gate"] == ( + "gate", + "shared", + ) + + +def test_qwen35_moe_reduces_routed_and_shared_outputs_together(): + class ReturnValue(torch.nn.Module): + def __init__(self, value): + super().__init__() + self.value = value + + def forward(self, *_): + return self.value + + hidden_states = torch.randn(2, 4) + local_output, shared_output = torch.randn_like(hidden_states), torch.randn_like(hidden_states) + gate_logits = torch.randn(2, 1) + block = Qwen35MoeSparseMoeBlock.__new__(Qwen35MoeSparseMoeBlock) + torch.nn.Module.__init__(block) + block.shared_expert = ReturnValue(shared_output) + block.gate = ReturnValue((torch.ones(2, 1), torch.zeros(2, 1, dtype=torch.int32), gate_logits)) + block.experts = ReturnValue(local_output) + block.parallel_context = Mock() + block.parallel_context.world.size = 2 + block.parallel_context.world_all_reduce.side_effect = lambda outputs: outputs + 1 + + with patch( + "sparsevllm.models.qwen3_5_moe.gated_shared_add", + side_effect=lambda routed, shared, gate: routed + shared * gate.sigmoid(), + ): + actual = block._forward_chunk(hidden_states) + + packed = block.parallel_context.world_all_reduce.call_args.args[0] + torch.testing.assert_close(packed, torch.stack((local_output, shared_output))) + torch.testing.assert_close( + actual, + local_output + 1 + (shared_output + 1) * gate_logits.sigmoid(), + ) + + +def test_qwen35_moe_skips_single_rank_output_packing(): + hidden_states = torch.randn(2, 4) + block = Qwen35MoeSparseMoeBlock.__new__(Qwen35MoeSparseMoeBlock) + torch.nn.Module.__init__(block) + block.shared_expert = Mock(return_value=hidden_states + 1) + block.gate = Mock( + return_value=( + torch.ones(2, 1), + torch.zeros(2, 1, dtype=torch.int32), + torch.zeros(2, 1), + ) + ) + block.experts = Mock(return_value=hidden_states + 2) + block.parallel_context = Mock() + block.parallel_context.world.size = 1 + + with patch( + "sparsevllm.models.qwen3_5_moe.gated_shared_add", + side_effect=lambda routed, shared, gate: routed + shared * gate.sigmoid(), + ): + actual = block._forward_chunk(hidden_states) + + block.parallel_context.world_all_reduce.assert_not_called() + torch.testing.assert_close( + actual, hidden_states + 2 + 0.5 * (hidden_states + 1) + ) + + class _ResidentAdmissionCache: def __init__(self): self.num_free_slots = 1_000_000 @@ -721,9 +904,11 @@ def stop_at_model_construction(_config): tensor_parallel_size=1, expert_parallel_size=1, data_parallel_size=1, + parallel_topology=ParallelTopology(1, 1, 1), uses_outer_tp_moe_layout=False, mlp_chunk_size=16384, hf_config=SimpleNamespace(model_type="qwen2", torch_dtype=torch.float32), + model_spec=resolve_model_spec("qwen2"), ) with ( patch.object(platforms, "_current_platform", platform), @@ -1016,7 +1201,7 @@ def test_qwen35_mixed_prefix_offload_allows_decode_graph(tmp_path): def test_qwen35_deltakv_requires_compatible_checkpoint_even_when_missing_allowed(tmp_path): - with pytest.raises(ValueError, match="DeltaKV for qwen3_5 requires"): + with pytest.raises(ValueError, match="DeltaKV for Qwen3.5 requires"): _make_config( tmp_path, vllm_sparse_method="deltakv", @@ -1042,7 +1227,7 @@ def test_qwen35_raw_config_fallback_when_transformers_autoconfig_is_unknown(tmp_ cfg = Config(model=str(tmp_path)) assert cfg.outer_hf_config.model_type == "qwen3_5" - assert cfg.hf_config.model_type == "qwen3_5" + assert cfg.hf_config.model_type == "qwen3_5_text" assert cfg.runtime_layout.num_kv_layers == 16 @@ -1064,7 +1249,7 @@ def test_qwen35_rejects_unquantized_fp16_checkpoint(tmp_path): outer_config.text_config.quantization_config = None with patch("sparsevllm.configs.runtime.AutoConfig.from_pretrained", return_value=outer_config): - with pytest.raises(NotImplementedError, match="require BF16 weights"): + with pytest.raises(ValueError, match="requires BF16 weights"): Config(model=str(tmp_path)) @@ -1073,7 +1258,7 @@ def test_qwen35_rejects_unsupported_quantization(tmp_path): outer_config.text_config.quantization_config = {"quant_method": "awq"} with patch("sparsevllm.configs.runtime.AutoConfig.from_pretrained", return_value=outer_config): - with pytest.raises(NotImplementedError, match="unquantized BF16 or block FP8"): + with pytest.raises(NotImplementedError, match="quant_method='awq'"): Config(model=str(tmp_path)) diff --git a/tests/test_qwen3_moe_compatibility.py b/tests/test_qwen3_moe_compatibility.py index 47aa506c..dd3f88e5 100644 --- a/tests/test_qwen3_moe_compatibility.py +++ b/tests/test_qwen3_moe_compatibility.py @@ -1,7 +1,10 @@ import pytest +from sparsevllm.distributed import ParallelMode, ParallelTopology from sparsevllm.method_registry import ( + DENSE_MODEL_COMPATIBILITY, MODEL_RUNTIME_COMPATIBILITY, + QWEN35_MOE_COMPATIBILITY, QWEN3_MOE_EP_COMPATIBILITY, QWEN3_MOE_TP_COMPATIBILITY, QWEN3_MOE_TP_EP_COMPATIBILITY, @@ -21,12 +24,23 @@ def _validate(method="", **overrides): "enable_prefix_caching": False, } values.update(overrides) + tp_size = values.pop("tensor_parallel_size") + ep_size = values.pop("expert_parallel_size") + dp_size = values.pop("data_parallel_size") + values["topology"] = ParallelTopology( + tp_size, + ep_size, + dp_size, + ParallelMode.OUTER_TP_MOE if tp_size > 1 else ParallelMode.STANDARD, + ) return validate_model_runtime_compatibility(**values) def test_qwen3_moe_registry_lists_only_v1_validated_combinations(): - assert MODEL_RUNTIME_COMPATIBILITY["qwen3_moe"] is QWEN3_MOE_EP_COMPATIBILITY - assert QWEN3_MOE_EP_COMPATIBILITY.parallel_mode == "ep_replicated_kv" + assert ( + MODEL_RUNTIME_COMPATIBILITY["qwen3_moe", ParallelMode.STANDARD] + is QWEN3_MOE_EP_COMPATIBILITY + ) assert QWEN3_MOE_EP_COMPATIBILITY.sparse_methods == { "", "streamingllm", @@ -59,6 +73,17 @@ def test_qwen3_moe_registry_lists_only_v1_validated_combinations(): } +def test_qwen35_moe_registry_accepts_vanilla_prefix_cache(): + assert validate_model_runtime_compatibility( + model_type="qwen3_5_moe", + sparse_method="", + topology=ParallelTopology(2, 2, 1, ParallelMode.OUTER_TP_MOE), + enforce_eager=True, + decode_cuda_graph=False, + enable_prefix_caching=True, + ) is QWEN35_MOE_COMPATIBILITY + + @pytest.mark.parametrize("method", sorted(QWEN3_MOE_EP_COMPATIBILITY.sparse_methods)) def test_qwen3_moe_registry_accepts_first_batch_sparse_methods(method): assert _validate(method) is QWEN3_MOE_EP_COMPATIBILITY @@ -79,9 +104,9 @@ def test_qwen3_moe_registry_rejects_unvalidated_prefix_cache_methods(method): def test_qwen3_moe_registry_rejects_conditional_and_out_of_scope_methods(): - with pytest.raises(NotImplementedError, match="steering asset"): + with pytest.raises(ValueError, match="validated methods"): _validate("skipkv") - with pytest.raises(NotImplementedError, match="not part of the validated"): + with pytest.raises(ValueError, match="validated methods"): _validate("deltakv") @@ -140,14 +165,20 @@ def test_qwen3_moe_registry_accepts_decode_cuda_graph(method): _validate(method, enforce_eager=False, decode_cuda_graph=True) is QWEN3_MOE_EP_COMPATIBILITY ) -def test_dense_models_do_not_inherit_qwen3_moe_compatibility(): +def test_dense_models_use_shared_runtime_compatibility(): assert validate_model_runtime_compatibility( model_type="qwen3", sparse_method="deltakv", - tensor_parallel_size=1, - expert_parallel_size=1, - data_parallel_size=1, + topology=ParallelTopology(1, 1, 1), enforce_eager=True, decode_cuda_graph=False, enable_prefix_caching=False, - ) is None + ) is DENSE_MODEL_COMPATIBILITY + + +def test_all_dense_architectures_are_registered(): + assert { + model_type + for model_type, mode in MODEL_RUNTIME_COMPATIBILITY + if mode is ParallelMode.STANDARD + } >= {"qwen2", "qwen3", "qwen3_5", "llama"} diff --git a/tests/test_snapkv_cache_budget.py b/tests/test_snapkv_cache_budget.py index 3ac34499..b8fe1b8e 100644 --- a/tests/test_snapkv_cache_budget.py +++ b/tests/test_snapkv_cache_budget.py @@ -5,6 +5,7 @@ import torch import sparsevllm.platforms as platforms +from sparsevllm.config import RuntimeLayout from sparsevllm.engine.cache_manager.base import CacheManager from sparsevllm.engine.cache_manager.rkv import RKVCacheManager from sparsevllm.engine.cache_manager.snapkv import ( @@ -37,7 +38,7 @@ def _manager_config(*, method: str, compression_interval: int = 1): head_dim=5, torch_dtype=torch.float32, ), - runtime_layout=None, + runtime_layout=RuntimeLayout.dense(2), max_model_len=5, max_num_batched_tokens=10, max_num_seqs_in_gpu=3, diff --git a/tests/test_sparsevllm_regression_grading.py b/tests/test_sparsevllm_regression_grading.py index c3dc692e..a2fe3317 100644 --- a/tests/test_sparsevllm_regression_grading.py +++ b/tests/test_sparsevllm_regression_grading.py @@ -41,6 +41,7 @@ main as run_suite_main, ) from benchmark.sparsevllm_regression.run_suite import _quality_command +from sparsevllm.config import RuntimeLayout from sparsevllm.engine.cache_manager.base import CacheManager from sparsevllm.distributed import ParallelContext, ParallelGroup from sparsevllm.method_registry import ( @@ -76,6 +77,7 @@ def __init__(self): ) config = types.SimpleNamespace( hf_config=hf_config, + runtime_layout=RuntimeLayout.dense(2), max_model_len=10, max_num_seqs_in_gpu=2, max_num_seqs_in_batch=2, @@ -191,7 +193,7 @@ def test_h2o_manifest_declares_supported_models_tp_runtime_matrix(self): method = manifest["methods"]["h2o"] self.assertEqual( method["supported_model_families"], - ["qwen2", "qwen3", "qwen3_moe", "qwen3_5", "llama", "minimax_m2"], + ["qwen2", "qwen3", "qwen3_moe", "qwen3_5", "qwen3_5_moe", "llama", "minimax_m2"], ) self.assertEqual( set(method["supported_model_families"]), diff --git a/tests/test_tp_rpc.py b/tests/test_tp_rpc.py index 50842041..c5a05448 100644 --- a/tests/test_tp_rpc.py +++ b/tests/test_tp_rpc.py @@ -19,6 +19,7 @@ TP_SHM_NAME_PREFIX, make_tp_shm_name, ) +from sparsevllm.operators import registry as operator_registry def test_write_shm_waits_until_worker_reads_command(): @@ -98,6 +99,20 @@ def test_prefix_offload_release_rpcs_use_failure_synchronized_world_path(): assert "free_slots_batch" in TP_RPC_STATUS_SYNC_METHODS +def test_operator_implementation_log_is_aligned_and_failure_synchronized(): + rank_zero = object.__new__(ModelRunner) + rank_zero.parallel_context = SimpleNamespace(world_rank=0) + rank_one = object.__new__(ModelRunner) + rank_one.parallel_context = SimpleNamespace(world_rank=1) + + with patch.object(operator_registry, "log_operator_implementations") as log_implementations: + ModelRunner.log_operator_implementations(rank_zero) + ModelRunner.log_operator_implementations(rank_one) + + assert "log_operator_implementations" in TP_RPC_STATUS_SYNC_METHODS + log_implementations.assert_called_once_with() + + def test_prefix_offload_release_rpc_surfaces_local_failure_after_status_sync(): for method_name, args in (("free_slots", (7,)), ("free_slots_batch", ([7, 9],))): runner = object.__new__(ModelRunner) diff --git a/tests/test_triton_moe.py b/tests/test_triton_moe.py index acd90195..a0c29f5a 100644 --- a/tests/test_triton_moe.py +++ b/tests/test_triton_moe.py @@ -4,7 +4,10 @@ import torch import torch.nn.functional as F +from sparsevllm.operators.gated_shared_add import gated_shared_add +from sparsevllm.triton_kernel.gate_up_swiglu import h20_gate_up_swiglu from sparsevllm.triton_kernel.moe import ( + _prepare_expert_assignment, fused_moe, fused_moe_gate_up_swiglu, moe_align_block_size, @@ -12,6 +15,45 @@ from sparsevllm.triton_kernel.moe_topk import topk_softmax +def _is_h20() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_name() == "NVIDIA H20" + + +@pytest.mark.skipif(not _is_h20(), reason="requires NVIDIA H20") +@pytest.mark.parametrize("intermediate_size", [256, 512]) +def test_h20_gate_up_swiglu_matches_torch(intermediate_size): + torch.manual_seed(0) + inputs = torch.randn(1, 2048, dtype=torch.bfloat16, device="cuda") + weight = 0.02 * torch.randn( + 2 * intermediate_size, + 2048, + dtype=torch.bfloat16, + device="cuda", + ) + projected = torch.nn.functional.linear(inputs, weight) + gate, up = projected.chunk(2, dim=-1) + expected = torch.nn.functional.silu(gate.float()) * up.float() + + actual = h20_gate_up_swiglu(inputs, weight) + + torch.testing.assert_close(actual.float(), expected, rtol=0.02, atol=0.01) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 1024]) +def test_gated_shared_add_matches_torch(num_tokens): + torch.manual_seed(num_tokens) + routed = torch.randn((num_tokens, 2048), device="cuda", dtype=torch.bfloat16) + shared = torch.randn_like(routed) + padded_gate = torch.randn((num_tokens, 257), device="cuda", dtype=torch.bfloat16) + gate_logits = padded_gate[:, -1:] + + actual = gated_shared_add(routed, shared, gate_logits) + expected = routed + torch.sigmoid(gate_logits) * shared + + torch.testing.assert_close(actual, expected, atol=0.03125, rtol=0.005) + + def _pytorch_topk_reference( logits: torch.Tensor, norm_topk_prob: bool, @@ -71,6 +113,22 @@ def test_moe_align_block_size_filters_ep_experts_and_pads_blocks(): assert sorted(sorted_ids[4:8].tolist()) == sorted([2, invalid, invalid, invalid]) +@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for Triton MoE tests.") +def test_naive_assignment_filters_ep_experts_in_one_kernel(): + alignment = _prepare_expert_assignment( + torch.tensor([[3, 12]], dtype=torch.int64, device="cuda"), + block_size=1, + num_experts=16, + local_expert_start=8, + local_expert_end=16, + ) + torch.cuda.synchronize() + + assert alignment.naive + assert alignment.expert_ids.tolist() == [-1, 4] + assert alignment.num_tokens_post_padded.item() == 2 + + @pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for Triton MoE tests.") def test_moe_alignment_covers_hotspot_and_empty_rank(dtype): @@ -138,6 +196,18 @@ def test_topk_softmax_accepts_any_valid_experts_for_ties(): ) +@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for Triton MoE tests.") +def test_topk_softmax_accepts_padded_row_stride(): + logits = torch.randn(3, 257, dtype=torch.bfloat16, device="cuda")[:, :256] + expected_weights, expected_ids = _pytorch_topk_reference(logits, True) + + weights, ids = topk_softmax(logits, top_k=8, norm_topk_prob=True) + torch.cuda.synchronize() + + assert torch.equal(ids, expected_ids.to(torch.int32)) + assert torch.allclose(weights.float(), expected_weights, atol=2e-2, rtol=2e-2) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for Triton MoE tests.") def test_topk_softmax_is_stable_for_extreme_finite_logits(): logits = torch.full((2, 128), -100, dtype=torch.bfloat16, device="cuda") diff --git a/tests/test_weight_loading.py b/tests/test_weight_loading.py index da142adc..e229ef1b 100644 --- a/tests/test_weight_loading.py +++ b/tests/test_weight_loading.py @@ -96,6 +96,39 @@ def record_skipped_weight( ) +class _CustomScaleLinear(nn.Module): + def __init__(self): + super().__init__() + self.quantized = True + self._quantized_weight_loaded = False + self.weight = nn.Parameter( + torch.empty(2, 2, dtype=torch.float8_e4m3fn), requires_grad=False + ) + self.register_buffer("weight_scale_inv", torch.empty(1, 1)) + + def load_quantized_weight(self, weight, scale, loaded_shard_id=None): + assert loaded_shard_id is None + self.weight.copy_(weight) + self.weight_scale_inv.copy_(scale.float()) + self._quantized_weight_loaded = True + + +class _CustomScaleModel(nn.Module): + def __init__(self): + super().__init__() + self.proj = _CustomScaleLinear() + + @staticmethod + def scale_key_for_weight(weight_key): + return weight_key.removesuffix(".weight") + ".scale" + + @staticmethod + def weight_key_for_scale(scale_key): + if not scale_key.endswith(".scale"): + return None + return scale_key.removesuffix(".scale") + ".weight" + + def _write_two_shards(path): left = torch.arange(4, dtype=torch.float32).reshape(2, 2) right = torch.arange(4, 8, dtype=torch.float32).reshape(2, 2) @@ -211,6 +244,21 @@ def test_load_model_keeps_remote_expert_tensors_as_metadata(tmp_path): ] +def test_load_model_supports_model_specific_scale_keys(tmp_path): + weight = torch.tensor([[1.0, 2.0], [3.0, 4.0]]).to(torch.float8_e4m3fn) + scale = torch.tensor([[0.5]], dtype=torch.float8_e8m0fnu) + save_file( + {"proj.weight": weight, "proj.scale": scale}, + tmp_path / "model.safetensors", + ) + model = _CustomScaleModel() + + loader.load_model(model, str(tmp_path), show_progress=False) + + torch.testing.assert_close(model.proj.weight.float(), weight.float()) + torch.testing.assert_close(model.proj.weight_scale_inv, scale.float()) + + def test_load_model_selects_all_files_for_local_checkpoint_rank(tmp_path): rank0_left = torch.full((2, 2), 10.0) rank0_right = torch.full((2, 2), 11.0)