From 956bf00015d693d21ca92beeb4082aa60f7c3b7b Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Mon, 10 Aug 2026 22:36:20 +0800 Subject: [PATCH 01/35] feat: support qwen3.6 moe bf16 --- README.md | 1 + README_zh.md | 1 + benchmark/long_bench/eval.py | 2 + benchmark/long_bench/pred.py | 173 +++++- benchmark/microbench.py | 85 +++ benchmark/runtime_validation.py | 25 + docs/en/features/supported-models.md | 8 + docs/zh/features/supported-models.md | 8 + .../validation/qwen36_compare_artifacts.py | 379 ++++++++++++ .../validation/qwen36_compare_longbench.py | 270 +++++++++ .../validation/qwen36_end_to_end_reference.py | 526 +++++++++++++++++ .../validation/qwen36_moe_bf16_reference.py | 367 ++++++++++++ .../validation/qwen36_summarize_microbench.py | 260 +++++++++ .../run_qwen36_microbench_matrix.py | 391 +++++++++++++ src/sparsevllm/configs/model.py | 187 +++++- src/sparsevllm/configs/runtime.py | 4 +- src/sparsevllm/engine/input_processor.py | 29 + src/sparsevllm/engine/llm_engine.py | 16 +- src/sparsevllm/engine/model_runner.py | 113 +++- .../engine/recurrent_state_manager.py | 29 +- src/sparsevllm/method_registry.py | 27 +- src/sparsevllm/models/qwen3_5.py | 33 +- src/sparsevllm/models/qwen3_5_moe.py | 542 ++++++++++++++++++ src/sparsevllm/operators/moe.py | 85 ++- src/sparsevllm/operators/moe_router.py | 156 +++++ src/sparsevllm/triton_kernel/moe_topk.py | 29 +- tests/test_input_processor.py | 54 ++ tests/test_longbench_deltakv_contracts.py | 134 ++++- tests/test_microbench_artifacts.py | 22 + tests/test_operator_providers.py | 119 +++- tests/test_prefill_schedule_policy.py | 6 +- tests/test_qwen35_mixed_runtime.py | 31 + tests/test_qwen35_moe.py | 403 +++++++++++++ tests/test_tp_rpc.py | 4 + tests/test_triton_moe.py | 9 +- 35 files changed, 4417 insertions(+), 111 deletions(-) create mode 100644 benchmark/runtime_validation.py create mode 100644 scripts/validation/qwen36_compare_artifacts.py create mode 100644 scripts/validation/qwen36_compare_longbench.py create mode 100644 scripts/validation/qwen36_end_to_end_reference.py create mode 100644 scripts/validation/qwen36_moe_bf16_reference.py create mode 100644 scripts/validation/qwen36_summarize_microbench.py create mode 100644 scripts/validation/run_qwen36_microbench_matrix.py create mode 100644 src/sparsevllm/engine/input_processor.py create mode 100644 src/sparsevllm/models/qwen3_5_moe.py create mode 100644 src/sparsevllm/operators/moe_router.py create mode 100644 tests/test_input_processor.py create mode 100644 tests/test_qwen35_moe.py diff --git a/README.md b/README.md index 05397d22..9350c932 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.6 MoE | ✅ (BF16 text-only) | | Llama 3 / 3.1 | ✅ | | MiniMax M2.7 | ✅ | diff --git a/README_zh.md b/README_zh.md index 306a5ffc..c18da47f 100644 --- a/README_zh.md +++ b/README_zh.md @@ -58,6 +58,7 @@ Sparse-vLLM 支持物理淘汰、逻辑掩码、查询感知选择和混合 KV | Qwen3 | ✅ | | Qwen3MoE | ✅ | | Qwen3.5 / Qwen3.6 | ✅ | +| Qwen3.6 MoE | ✅（仅 BF16 纯文本） | | Llama 3 / 3.1 | ✅ | | MiniMax M2.7 | ✅ | diff --git a/benchmark/long_bench/eval.py b/benchmark/long_bench/eval.py index e4698f02..1126ae71 100644 --- a/benchmark/long_bench/eval.py +++ b/benchmark/long_bench/eval.py @@ -327,5 +327,7 @@ def aggregate_category_scores(task_scores): json.dump(scores, f, ensure_ascii=False, indent=4) with open(os.path.join(path, "metrics.json"), "w") as f: json.dump(scores, f, ensure_ascii=False, indent=4) + with open(os.path.join(path, "aggregate_metrics.json"), "w") as f: + json.dump(scores, f, ensure_ascii=False, indent=4) if failed_tasks: raise SystemExit(1) diff --git a/benchmark/long_bench/pred.py b/benchmark/long_bench/pred.py index 59fc46fe..9e3714ba 100644 --- a/benchmark/long_bench/pred.py +++ b/benchmark/long_bench/pred.py @@ -1,5 +1,6 @@ import os import json +import hashlib import sys import subprocess import re @@ -23,6 +24,7 @@ from transformers import AutoTokenizer, GenerationConfig import torch.distributed as dist from benchmark.model_adapters.sparsevllm import get_sparsevllm_generate_api +from benchmark.runtime_validation import collect_worker_runtime_status from datetime import datetime BASE_PATH = os.getenv("SPARSEVLLM_OUTPUT_DIR", str(REPO_ROOT / "outputs")) @@ -38,6 +40,30 @@ } +def _sha256(path: str | os.PathLike[str]) -> str: + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _sha256_text(value: str) -> str: + return hashlib.sha256(value.encode("utf-8")).hexdigest() + + +def _git_value(*args: str) -> str | None: + result = subprocess.run( + ["git", *args], + cwd=REPO_ROOT, + text=True, + capture_output=True, + check=False, + ) + value = result.stdout.strip() + return value or None + + def get_longbench_data_path(dataset, use_longbench_e): if not DATA_PREFIX_PATH: raise FileNotFoundError( @@ -145,6 +171,7 @@ def _artifact_paths(out_root: str) -> dict[str, str]: "raw": os.path.join(out_root, "raw_outputs.jsonl"), "parsed": os.path.join(out_root, "parsed_outputs.jsonl"), "sample": os.path.join(out_root, "sample_results.jsonl"), + "per_sample": os.path.join(out_root, "per_sample_results.jsonl"), } @@ -161,28 +188,23 @@ def _write_decode_cuda_graph_status( "cannot verify decode CUDA graph execution." ) - runner = getattr(llm, "model_runner", None) - graph_runner = getattr(runner, "decode_cuda_graph_runner", None) - graph_states = ( - getattr(graph_runner, "_graphs", {}) - if graph_runner is not None - else {} - ) - graph_count = sum( - getattr(state, "graph", None) is not None - for state in graph_states.values() + statuses = collect_worker_runtime_status(llm) + if not statuses: + raise RuntimeError("Runtime validation returned no worker status records.") + configured_flags = [ + bool(status.get("decode_cuda_graph_configured")) for status in statuses + ] + configured = all(configured_flags) + expected = bool(getattr(llm.config, "decode_cuda_graph", False)) + all_active = all( + bool(status.get("decode_cuda_graph_active")) for status in statuses ) graph_status = { - "rank": int(rank), - "configured": bool( - getattr(getattr(llm, "config", None), "decode_cuda_graph", False) - ), - "runner_initialized": graph_runner is not None, - "state_count": int(len(graph_states)), - "graph_count": int(graph_count), - "active": bool(graph_count > 0), - "last_state_key": str(getattr(graph_runner, "last_state_key", None)), - "state_keys": [str(key) for key in graph_states], + "launcher_rank": int(rank), + "expected": expected, + "configured_on_all_workers": configured, + "active_on_all_workers": all_active, + "workers": statuses, } status_path = os.path.join( out_root, @@ -191,6 +213,16 @@ def _write_decode_cuda_graph_status( with open(status_path, "w", encoding="utf-8") as handle: json.dump(graph_status, handle, ensure_ascii=False, indent=2) handle.write("\n") + if expected and not configured: + raise RuntimeError( + "Decode CUDA Graph was requested but is not configured on every worker: " + f"{statuses!r}." + ) + if configured and not all_active: + raise RuntimeError( + "Decode CUDA Graph was requested but was not active on every model worker: " + f"{statuses!r}." + ) return graph_status @@ -234,6 +266,8 @@ def _write_sample_record( "source_idx", "status", "prompt_tokens", + "rendered_prompt", + "rendered_prompt_sha256", "raw_pred", "error", "traceback", @@ -256,6 +290,7 @@ def _write_sample_record( _append_jsonl(paths["raw"], raw_record) _append_jsonl(paths["parsed"], parsed_record) _append_jsonl(paths["sample"], record) + _append_jsonl(paths["per_sample"], record) # Keep the historical per-task files for benchmark/long_bench/eval.py. task_record = { @@ -374,6 +409,10 @@ def get_pred(rank, data, dataset_info, args, model, tokenizer, model_max_length, prompt_tokens=prompt_tokens, ) ) + prepared_records[-1]["rendered_prompt"] = prompt + prepared_records[-1]["rendered_prompt_sha256"] = _sha256_text( + prompt + ) except Exception as exc: record = _sample_base_record( dataset=dataset, @@ -491,8 +530,28 @@ def get_pred(rank, data, dataset_info, args, model, tokenizer, model_max_length, def worker(rank, world_size, datasets, dataset2prompt, dataset2maxlen, args, out_root, max_length_limit): - seed_everything(42) + seed_everything(args.seed) model, tokenizer, model_max_length, eos_token_ids = load_model_and_tokenizer(rank, args) + if rank == 0: + tokenizer_runtime = { + "tokenizer_class": type(tokenizer).__name__, + "tokenizer_path": args.tokenizer_path or args.model_path, + "chat_template": tokenizer.chat_template, + "bos_token": tokenizer.bos_token, + "bos_token_id": tokenizer.bos_token_id, + "eos_token": tokenizer.eos_token, + "eos_token_id": tokenizer.eos_token_id, + "effective_eos_token_ids": eos_token_ids, + "no_chat_template": args.no_chat_template, + "thinking_mode": args.thinking_mode, + } + with open( + os.path.join(out_root, "tokenizer_runtime.json"), + "w", + encoding="utf-8", + ) as handle: + json.dump(tokenizer_runtime, handle, ensure_ascii=False, indent=2) + handle.write("\n") for dataset in datasets: data_path = get_longbench_data_path(dataset, args.e) @@ -665,6 +724,7 @@ def parse_args(): parser.add_argument("--min_prompt_tokens", type=int, default=None) parser.add_argument("--samples_per_task", type=int, default=20) parser.add_argument("--min_required_samples", type=int, default=5) + parser.add_argument("--seed", type=int, default=20260810) parser.add_argument("--worker_rank", type=int, default=-1) parser.add_argument("--worker_world_size", type=int, default=1) parser.add_argument("--output_root", type=str, default=None) @@ -707,7 +767,13 @@ def parse_args(): for dataset in datasets: with open(os.path.join(out_root, f"{dataset}.jsonl"), 'w') as f: pass - for artifact in ("raw_outputs.jsonl", "parsed_outputs.jsonl", "sample_results.jsonl", "longbench_mini_selection.jsonl"): + for artifact in ( + "raw_outputs.jsonl", + "parsed_outputs.jsonl", + "sample_results.jsonl", + "per_sample_results.jsonl", + "longbench_mini_selection.jsonl", + ): with open(os.path.join(out_root, artifact), "w", encoding="utf-8") as f: pass @@ -715,16 +781,79 @@ def parse_args(): args.max_model_len = max_length_limit if args.worker_rank < 0: + dataset_files = [] + for dataset in datasets: + data_path = Path(get_longbench_data_path(dataset, args.e)).resolve() + dataset_files.append( + { + "dataset": dataset, + "path": str(data_path), + "size_bytes": data_path.stat().st_size, + "sha256": _sha256(data_path), + } + ) + model_path = Path(args.model_path).resolve() + model_files = {} + for name in ("config.json", "model.safetensors.index.json"): + path = model_path / name + if not path.is_file(): + raise FileNotFoundError(f"Required model metadata file is missing: {path}") + model_files[name] = {"path": str(path), "sha256": _sha256(path)} + tokenizer_files = {} + for name in ( + "tokenizer.json", + "tokenizer_config.json", + "generation_config.json", + "chat_template.jinja", + ): + path = model_path / name + if path.is_file(): + tokenizer_files[name] = { + "path": str(path), + "sha256": _sha256(path), + } + prompt_config_path = REPO_ROOT / "benchmark/long_bench/config/dataset2prompt.json" + maxlen_config_path = REPO_ROOT / "benchmark/long_bench/config/dataset2maxlen.json" resolved_config = { + "created_at": datetime.now().isoformat(timespec="seconds"), + "command": " ".join([sys.executable, *sys.argv]), + "git_commit": _git_value("rev-parse", "HEAD"), + "git_branch": _git_value("branch", "--show-current"), + "git_dirty": bool(_git_value("status", "--porcelain")), "model": args.model, "model_path": args.model_path, + "model_files": model_files, + "tokenizer_files": tokenizer_files, "tokenizer_path": args.tokenizer_path or args.model_path, "backend": "sparsevllm", + "provider_env": { + key: os.environ.get(key, "auto") + for key in ( + "SPARSEVLLM_MOE_PROVIDER", + "SPARSEVLLM_MOE_ROUTER_PROVIDER", + ) + }, "sparse_method": args.sparse_method, "deltakv_checkpoint_path": args.deltakv_checkpoint_path, "datasets": datasets, "longbench_data_root": DATA_PREFIX_PATH, + "dataset_files": dataset_files, + "prompt_config": { + "path": str(prompt_config_path), + "sha256": _sha256(prompt_config_path), + "selected_formats": { + dataset: dataset2prompt[dataset] for dataset in datasets + }, + }, + "maxlen_config": { + "path": str(maxlen_config_path), + "sha256": _sha256(maxlen_config_path), + "selected_values": { + dataset: dataset2maxlen[dataset] for dataset in datasets + }, + }, "max_model_len": args.max_model_len, + "seed": args.seed, "decoding": { "temperature": args.temperature, "top_p": args.top_p, diff --git a/benchmark/microbench.py b/benchmark/microbench.py index ffa5bf76..610df79e 100644 --- a/benchmark/microbench.py +++ b/benchmark/microbench.py @@ -27,6 +27,7 @@ is_tp_decode_cuda_graph_supported, normalize_sparse_method, ) +from benchmark.runtime_validation import collect_worker_runtime_status DEFAULT_ALL_CHUNKED_PREFILL_SIZE = 96 * 1024 @@ -156,6 +157,8 @@ def _selected_env_snapshot() -> dict[str, str]: "SPARSEVLLM_LONG_PREFILL_OFFLOAD_MIN_TOKENS", "SPARSEVLLM_RAWKV_BUFFER_MODE", "SPARSEVLLM_RAWKV_PREFETCH", + "SPARSEVLLM_MOE_ROUTER_PROVIDER", + "SPARSEVLLM_MOE_PROVIDER", ] return {key: os.environ[key] for key in keys if key in os.environ} @@ -276,6 +279,10 @@ def _artifact_records(args, rows: list[dict[str, Any]]) -> list[dict[str, Any]]: "synchronize_step_timing", bool(getattr(args, "synchronize_step_timing", False)), ) + record.setdefault( + "warmup_output_len", + int(getattr(args, "warmup_output_len", 0) or 0), + ) if "prefill_tp" in row: record.setdefault("prefill_tok_s", row["prefill_tp"]) if "decode_tp" in row: @@ -286,6 +293,8 @@ def _artifact_records(args, rows: list[dict[str, Any]]) -> list[dict[str, Any]]: record.setdefault("itl_ms", row["itl"]) if "mem" in row: record.setdefault("peak_memory_gb", row["mem"]) + if "end_to_end_tp" in row: + record.setdefault("end_to_end_tok_s", row["end_to_end_tp"]) records.append(record) return records @@ -311,9 +320,11 @@ def _write_output_dir(args, rows: list[dict[str, Any]]) -> None: "output_len": int(args.output_len), "temperature": float(args.temperature), "top_p": float(args.top_p), + "seed": int(getattr(args, "seed", 20260810)), "synchronize_step_timing": bool( getattr(args, "synchronize_step_timing", False) ), + "warmup_output_len": int(getattr(args, "warmup_output_len", 0) or 0), "hyper_params": args.hyper_params_dict, "env": _selected_env_snapshot(), } @@ -391,6 +402,10 @@ def _decode_cuda_graph_status(llm) -> dict[str, Any]: } +def _worker_runtime_status(llm) -> list[dict[str, Any]]: + return collect_worker_runtime_status(llm) + + def _jsonable_config_value(value: Any) -> Any: if value is None or isinstance(value, (str, int, float, bool)): return value @@ -472,6 +487,9 @@ def _finished_outputs_have_tokens(finished_outputs) -> bool: def benchmark_task(method, length, bs, args, results_dict): + seed = int(getattr(args, "seed", 20260810)) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) # 为每个子进程重置显存统计 torch.cuda.reset_peak_memory_stats() torch.cuda.empty_cache() @@ -541,6 +559,41 @@ def benchmark_task(method, length, bs, args, results_dict): } llm = LLM(args.model_path, **engine_kwargs) resolved_engine_config = _resolved_engine_config(llm) + warmup_output_len = int(getattr(args, "warmup_output_len", 0) or 0) + if warmup_output_len < 0: + raise ValueError("warmup_output_len must be non-negative.") + if warmup_output_len: + warmup_prompt_len = min(int(length), 1024) + warmup_prompts = [[100] * warmup_prompt_len for _ in range(bs)] + warmup_sampling = [ + SamplingParams( + temperature=0.0, + top_p=1.0, + ignore_eos=True, + max_tokens=warmup_output_len, + ) + for _ in range(bs) + ] + warmup_outputs = llm.generate( + warmup_prompts, + warmup_sampling, + use_tqdm=False, + ) + if len(warmup_outputs) != bs or not llm.is_finished(): + raise RuntimeError( + "Microbenchmark warmup did not finish every request: " + f"expected={bs}, outputs={len(warmup_outputs)}, " + f"engine_finished={llm.is_finished()}." + ) + if bool(base_hyper_params.get("decode_cuda_graph")) and not all( + bool(status.get("decode_cuda_graph_active")) + for status in _worker_runtime_status(llm) + ): + raise RuntimeError( + "Decode CUDA Graph was requested but did not activate during warmup." + ) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() prefix_cache_stats_before = _cache_stats(llm) prompt_token_ids = [[100] * length for _ in range(bs)] @@ -686,8 +739,25 @@ def add_wave(max_new_requests: int): t_end = perf_counter() duration = t_end - t_start + end_to_end_tokens = int(prefill_tokens + decode_tokens) + end_to_end_tp = end_to_end_tokens / duration if duration > 0 else 0.0 peak_mem = get_peak_memory() graph_status = _decode_cuda_graph_status(llm) + worker_runtime_status = _worker_runtime_status(llm) + primary_worker_status = worker_runtime_status[0] + provider_status = { + key: primary_worker_status[key] + for key in ("moe_expert_provider", "moe_router_provider") + if key in primary_worker_status + } + if bool(base_hyper_params.get("decode_cuda_graph")) and not all( + bool(status.get("decode_cuda_graph_active")) + for status in worker_runtime_status + ): + raise RuntimeError( + "Decode CUDA Graph was requested but was not active on every " + f"worker: {worker_runtime_status!r}." + ) prefix_cache_stats_after = _cache_stats(llm) prefix_cache_stats_delta = _numeric_delta(prefix_cache_stats_before, prefix_cache_stats_after) observed_prefix_hit_tokens = int(sum(prefix_hits_by_seq_id.values())) @@ -757,6 +827,9 @@ def add_wave(max_new_requests: int): "itl": avg_itl, "avg_bs": avg_active_bs, "mem": peak_mem, + "duration_s": duration, + "end_to_end_tokens": end_to_end_tokens, + "end_to_end_tp": end_to_end_tp, "has_queued": has_queued, "full_admission_reached": full_admission_reached, "impossible_full_admission": impossible_full_admission, @@ -773,6 +846,8 @@ def add_wave(max_new_requests: int): "scheduler_recompute_replays": recompute_replays, "decode_cuda_graph_expected": bool(base_hyper_params.get("decode_cuda_graph")), **graph_status, + **provider_status, + "worker_runtime_status": worker_runtime_status, "prefix_cache_required": bool(getattr(args, "require_prefix_cache_hit", False)), "prefix_cache_stats_before": prefix_cache_stats_before, "prefix_cache_stats_after": prefix_cache_stats_after, @@ -843,6 +918,7 @@ def main(): default=1.0, help="Nucleus sampling top-p. Only used when temperature > 0.", ) + parser.add_argument("--seed", type=int, default=20260810) parser.add_argument( "--admission_wave_size", type=int, @@ -869,6 +945,15 @@ def main(): "so post-sparse work is attributed to the step that launched it." ), ) + parser.add_argument( + "--warmup_output_len", + type=int, + default=0, + help=( + "Run an unmeasured same-batch warmup before each case; a positive value " + "also forces decode CUDA Graph capture before timed steps." + ), + ) parser.add_argument( "--wave_decode_gap_steps", type=int, diff --git a/benchmark/runtime_validation.py b/benchmark/runtime_validation.py new file mode 100644 index 00000000..d6263017 --- /dev/null +++ b/benchmark/runtime_validation.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from typing import Any + + +def collect_worker_runtime_status(llm) -> list[dict[str, Any]]: + """Collect worker diagnostics without extending the public engine API.""" + model_runner = getattr(llm, "model_runner", None) + call = getattr(model_runner, "call", None) + if not callable(call): + raise RuntimeError( + "Sparse-VLLM runtime validation requires model_runner.call()." + ) + statuses = call("runtime_diagnostic_status") + expected = int(getattr(getattr(llm, "config", None), "world_size", 1)) + if ( + not isinstance(statuses, list) + or len(statuses) != expected + or not all(isinstance(status, dict) for status in statuses) + ): + raise RuntimeError( + "Runtime validation must return one status object per model worker: " + f"expected={expected}, got={statuses!r}." + ) + return statuses diff --git a/docs/en/features/supported-models.md b/docs/en/features/supported-models.md index 7bb3942a..748c0ffd 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 only | ✅ | 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 | ✅ | @@ -31,6 +32,12 @@ dimension must be divisible by `T / E`. Qwen3MoE outer TP requires a BF16 model dtype; FP16 Qwen3MoE checkpoints are limited to `TP=1`. When `TP=1`, the existing EP layout uses world size `E`. +Qwen3.6 MoE always uses the outer-TP layout: attention and Gated DeltaNet TP +are `T`, MoE EP is `E`, MoE TP is `T / E`, and world size is `T`. It requires +`DP=1`, `T % E == 0`, and BF16 language-model weights. The first release is a +text-only CausalLM runtime, rejects image/video and MTP inputs, supports only +Vanilla KV runtime, and captures decode (not prefill) with CUDA Graph. + 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`. @@ -43,6 +50,7 @@ normalized internally to `model_type=qwen3_5`. | Qwen3 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | Compressor required² | | Qwen3MoE | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Qwen3.5 / Qwen3.6 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | Matched checkpoint³ | +| Qwen3.6 MoE | ✅ | — | — | — | — | — | — | — | — | — | | Llama 3 / 3.1 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | Selected checkpoint¹ | Compressor required² | | MiniMax M2.7 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | diff --git a/docs/zh/features/supported-models.md b/docs/zh/features/supported-models.md index d273ac8d..eab75258 100644 --- a/docs/zh/features/supported-models.md +++ b/docs/zh/features/supported-models.md @@ -12,6 +12,7 @@ | 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 | ✅ | 仅支持 1 | ✅ | | Llama 3 / 3.1 | `llama` | BF16 / FP16 | ✅ | 仅支持 1 | 仅支持 1 | | MiniMax M2.7 | `minimax_m2` | 块级 FP8，非量化权重使用 BF16 | ✅ | 仅支持 1 | ✅ | @@ -24,6 +25,12 @@ size 为 `T`。该布局要求 `DP=1` 且 `T % E == 0`；专家数量必须能 TP 要求模型 dtype 为 BF16；FP16 Qwen3MoE checkpoint 仅支持 `TP=1`。当 `TP=1` 时，原有 EP 布局的 world size 为 `E`。 +Qwen3.6 MoE 始终使用 outer-TP 布局：attention 与 Gated DeltaNet TP 为 +`T`、MoE EP 为 `E`、MoE TP 为 `T / E`，world size 为 `T`。该模型要求 +`DP=1`、`T % E == 0` 且语言模型权重为 BF16。首版仅支持纯文本 CausalLM， +明确拒绝 image/video 与 MTP 输入，只支持 Vanilla KV runtime；CUDA Graph +仅覆盖 decode，不覆盖 prefill。 + 块级 FP8 要求使用 E4M3 权重、动态激活量化以及 `128 x 128` 的权重块大小。Qwen3.5 和 Qwen3.6 的配置在内部统一规范为 `model_type=qwen3_5`。 ## 稀疏方法支持 @@ -34,6 +41,7 @@ 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 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | diff --git a/scripts/validation/qwen36_compare_artifacts.py b/scripts/validation/qwen36_compare_artifacts.py new file mode 100644 index 00000000..1d6c112d --- /dev/null +++ b/scripts/validation/qwen36_compare_artifacts.py @@ -0,0 +1,379 @@ +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from datetime import datetime +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F + + +THRESHOLDS = { + "forced_logits_max_abs": 4.0, + "forced_logits_mean_abs": 0.6, + "near_tie_margin": 0.25, + "decoder_layer_cosine": 0.75, + "routing_layer_overlap": 0.90, + "routing_mean_overlap": 0.95, + "graph_eager_logits_max_abs": 0.0, +} + + +def _read_json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise TypeError(f"Expected a JSON object in {path}, got {type(value).__name__}.") + return value + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + rows = [] + for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + try: + row = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError(f"Invalid JSONL at {path}:{line_no}: {exc}") from exc + if not isinstance(row, dict): + raise TypeError(f"Expected an object at {path}:{line_no}.") + rows.append(row) + return rows + + +def _load_tensor_list(path: Path) -> list[Any]: + value = torch.load(path, map_location="cpu", weights_only=False) + if not isinstance(value, list): + raise TypeError(f"Expected a tensor list in {path}, got {type(value).__name__}.") + return value + + +def _require_run(path: Path) -> None: + required = ( + "run_info.json", + "runtime_status.json", + "per_sample_results.jsonl", + "aggregate_metrics.json", + "raw_logits.pt", + ) + missing = [name for name in required if not (path / name).is_file()] + if missing: + raise FileNotFoundError(f"Validation run {path} is missing artifacts: {missing}.") + aggregate = _read_json(path / "aggregate_metrics.json") + if aggregate.get("status") != "success": + raise RuntimeError(f"Validation run {path} is not successful: {aggregate!r}.") + + +def _tensor_metrics(actual: torch.Tensor, reference: torch.Tensor) -> dict[str, float]: + if actual.shape != reference.shape: + raise ValueError( + f"Tensor shape mismatch: actual={tuple(actual.shape)}, reference={tuple(reference.shape)}." + ) + actual_fp32 = actual.float() + reference_fp32 = reference.float() + difference = (actual_fp32 - reference_fp32).abs() + return { + "max_abs": float(difference.max().item()), + "mean_abs": float(difference.mean().item()), + "cosine": float( + F.cosine_similarity(actual_fp32.flatten(), reference_fp32.flatten(), dim=0).item() + ), + } + + +def _output_ids(path: Path) -> list[list[int]]: + rows = _read_jsonl(path / "per_sample_results.jsonl") + if any(row.get("status") != "success" for row in rows): + raise RuntimeError(f"Non-success sample found in {path}.") + return [[int(token) for token in row["output_token_ids"]] for row in rows] + + +def _graph_records( + label: str, + eager_dir: Path, + graph_dir: Path, +) -> list[dict[str, Any]]: + eager_logits = _load_tensor_list(eager_dir / "raw_logits.pt") + graph_logits = _load_tensor_list(graph_dir / "raw_logits.pt") + if len(eager_logits) != len(graph_logits): + raise ValueError(f"{label} eager/Graph sample counts differ.") + max_abs = max( + _tensor_metrics(graph, eager)["max_abs"] + for eager, graph in zip(eager_logits, graph_logits) + ) + token_ids_equal = _output_ids(eager_dir) == _output_ids(graph_dir) + runtime = _read_json(graph_dir / "runtime_status.json") + workers = runtime.get("worker_runtime_status") + if not isinstance(workers, list) or not workers: + raise RuntimeError(f"{graph_dir} has no worker_runtime_status records.") + all_graph_active = all( + worker.get("decode_cuda_graph_configured") is True + and worker.get("decode_cuda_graph_active") is True + for worker in workers + ) + providers = { + (worker.get("moe_expert_provider"), worker.get("moe_router_provider")) + for worker in workers + } + status = ( + "success" + if max_abs <= THRESHOLDS["graph_eager_logits_max_abs"] + and token_ids_equal + and all_graph_active + and providers == {("triton", "triton")} + else "metric_failed" + ) + return [ + { + "check": "graph_eager_equivalence", + "topology": label, + "status": status, + "max_abs": max_abs, + "token_ids_equal": token_ids_equal, + "all_graph_active": all_graph_active, + "worker_count": len(workers), + "providers": sorted([list(pair) for pair in providers]), + } + ] + + +def _forced_reference_records( + label: str, + reference_dir: Path, + actual_dir: Path, +) -> list[dict[str, Any]]: + reference_logits = _load_tensor_list(reference_dir / "raw_logits.pt") + actual_logits = _load_tensor_list(actual_dir / "raw_logits.pt") + reference_hidden = _load_tensor_list( + reference_dir / "raw_cached_hidden_states.pt" + ) + actual_hidden = _load_tensor_list(actual_dir / "raw_cached_hidden_states.pt") + if not ( + len(reference_logits) + == len(actual_logits) + == len(reference_hidden) + == len(actual_hidden) + ): + raise ValueError(f"{label} forced-reference sample counts differ.") + + records: list[dict[str, Any]] = [] + for sample_idx, (reference, actual) in enumerate( + zip(reference_logits, actual_logits) + ): + metrics = _tensor_metrics(actual, reference) + reference_top2 = torch.topk(reference.float(), k=2) + actual_top1 = int(torch.argmax(actual).item()) + reference_top1 = int(reference_top2.indices[0].item()) + reference_margin = float( + (reference_top2.values[0] - reference_top2.values[1]).item() + ) + top1_acceptable = actual_top1 == reference_top1 or ( + reference_margin <= THRESHOLDS["near_tie_margin"] + and actual_top1 in {int(index) for index in reference_top2.indices.tolist()} + ) + status = ( + "success" + if metrics["max_abs"] <= THRESHOLDS["forced_logits_max_abs"] + and metrics["mean_abs"] <= THRESHOLDS["forced_logits_mean_abs"] + and top1_acceptable + else "metric_failed" + ) + records.append( + { + "check": "forced_prefix_logits", + "topology": label, + "sample_idx": sample_idx, + "status": status, + **metrics, + "reference_top1": reference_top1, + "actual_top1": actual_top1, + "reference_top1_margin": reference_margin, + "top1_acceptable": top1_acceptable, + } + ) + + reference_layers = reference_hidden[sample_idx] + actual_layers = actual_hidden[sample_idx] + if set(reference_layers) != set(actual_layers): + raise ValueError(f"{label} sample {sample_idx} hidden layer sets differ.") + for layer_idx in sorted(reference_layers): + layer_metrics = _tensor_metrics( + actual_layers[layer_idx], reference_layers[layer_idx] + ) + records.append( + { + "check": "decoder_layer_output", + "topology": label, + "sample_idx": sample_idx, + "layer_idx": int(layer_idx), + "layer_kind": ( + "embedding" + if layer_idx == -1 + else "final_norm" + if layer_idx == 40 + else "full_attention" + if layer_idx % 4 == 3 + else "gated_deltanet" + ), + "status": ( + "success" + if layer_metrics["cosine"] + >= THRESHOLDS["decoder_layer_cosine"] + else "metric_failed" + ), + **layer_metrics, + } + ) + return records + + +def _routing_records( + label: str, + reference_dir: Path, + actual_dir: Path, +) -> list[dict[str, Any]]: + reference = _load_tensor_list(reference_dir / "raw_moe_states.pt") + actual = _load_tensor_list(actual_dir / "raw_moe_states.pt") + if len(reference) != len(actual): + raise ValueError(f"{label} routing sample counts differ.") + records = [] + overlaps = [] + for sample_idx, (reference_layers, actual_layers) in enumerate( + zip(reference, actual) + ): + if set(reference_layers) != set(actual_layers): + raise ValueError(f"{label} sample {sample_idx} MoE layer sets differ.") + for layer_idx in sorted(reference_layers): + reference_ids = reference_layers[layer_idx]["topk_ids"] + actual_ids = actual_layers[layer_idx]["topk_ids"] + if reference_ids.shape != actual_ids.shape: + raise ValueError(f"{label} sample {sample_idx} layer {layer_idx} shape differs.") + row_overlaps = [ + len(set(left.tolist()) & set(right.tolist())) / reference_ids.shape[1] + for left, right in zip(reference_ids, actual_ids) + ] + overlap = float(sum(row_overlaps) / len(row_overlaps)) + overlaps.append(overlap) + records.append( + { + "check": "cross_topology_routing", + "topology": label, + "sample_idx": sample_idx, + "layer_idx": int(layer_idx), + "status": ( + "success" + if overlap >= THRESHOLDS["routing_layer_overlap"] + else "metric_failed" + ), + "topk_set_overlap": overlap, + "ordered_ids_equal": bool(torch.equal(reference_ids, actual_ids)), + } + ) + mean_overlap = float(sum(overlaps) / len(overlaps)) + records.append( + { + "check": "cross_topology_routing_aggregate", + "topology": label, + "status": ( + "success" + if mean_overlap >= THRESHOLDS["routing_mean_overlap"] + else "metric_failed" + ), + "mean_topk_set_overlap": mean_overlap, + "min_layer_topk_set_overlap": min(overlaps), + } + ) + return records + + +def _write_json(path: Path, value: Any) -> None: + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Compare fixed Qwen3.6 MoE correctness artifacts." + ) + parser.add_argument("--transformers", type=Path, required=True) + for name in ( + "single-eager", + "single-graph", + "tp-eager", + "tp-graph", + "tp-ep-eager", + "tp-ep-graph", + "forced-single", + "forced-tp", + "forced-tp-ep", + ): + parser.add_argument(f"--{name}", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + args = parser.parse_args() + + sources = { + key: Path(value).resolve() + for key, value in vars(args).items() + if key != "output_dir" + } + for path in sources.values(): + _require_run(path) + args.output_dir = args.output_dir.resolve() + args.output_dir.mkdir(parents=True, exist_ok=False) + + records = [] + records += _graph_records("single", sources["single_eager"], sources["single_graph"]) + records += _graph_records("pure_tp", sources["tp_eager"], sources["tp_graph"]) + records += _graph_records("tp_ep", sources["tp_ep_eager"], sources["tp_ep_graph"]) + for label, key in ( + ("single", "forced_single"), + ("pure_tp", "forced_tp"), + ("tp_ep", "forced_tp_ep"), + ): + records += _forced_reference_records( + label, sources["transformers"], sources[key] + ) + records += _routing_records( + "single_vs_pure_tp", sources["forced_single"], sources["forced_tp"] + ) + records += _routing_records( + "single_vs_tp_ep", sources["forced_single"], sources["forced_tp_ep"] + ) + + failed = [record for record in records if record["status"] != "success"] + aggregate = { + "status": "success" if not failed else "metric_failed", + "num_checks": len(records), + "success_checks": len(records) - len(failed), + "failed_checks": len(failed), + "thresholds": THRESHOLDS, + } + run_info = { + "created_at": datetime.now().isoformat(timespec="seconds"), + "command": " ".join(sys.argv), + "git_commit": subprocess.run( + ["git", "rev-parse", "HEAD"], text=True, capture_output=True, check=True + ).stdout.strip(), + "sources": {key: str(value) for key, value in sources.items()}, + "thresholds": THRESHOLDS, + } + _write_json(args.output_dir / "run_info.json", run_info) + _write_json(args.output_dir / "raw_outputs.json", run_info["sources"]) + _write_json(args.output_dir / "parsed_outputs.json", {"checks": records}) + with (args.output_dir / "per_sample_results.jsonl").open( + "w", encoding="utf-8" + ) as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + _write_json(args.output_dir / "aggregate_metrics.json", aggregate) + print(json.dumps(aggregate, ensure_ascii=False, indent=2)) + if failed: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/scripts/validation/qwen36_compare_longbench.py b/scripts/validation/qwen36_compare_longbench.py new file mode 100644 index 00000000..f7476826 --- /dev/null +++ b/scripts/validation/qwen36_compare_longbench.py @@ -0,0 +1,270 @@ +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from datetime import datetime +from pathlib import Path +from typing import Any + + +TASK_SCORE_MAX_DROP = 20.0 +MEAN_SCORE_MAX_DROP = 3.0 + + +def _read_json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise TypeError(f"Expected a JSON object in {path}.") + return value + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + rows = [] + for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + value = json.loads(line) + if not isinstance(value, dict): + raise TypeError(f"Expected an object at {path}:{line_no}.") + rows.append(value) + return rows + + +def _load_run(path: Path) -> dict[str, Any]: + required = ( + "resolved_config.json", + "raw_outputs.jsonl", + "parsed_outputs.jsonl", + "per_sample_results.jsonl", + "aggregate_metrics.json", + "decode_cuda_graph_status_rank0.json", + "tokenizer_runtime.json", + ) + missing = [name for name in required if not (path / name).is_file()] + if missing: + raise FileNotFoundError(f"LongBench run {path} is missing {missing}.") + metrics = _read_json(path / "aggregate_metrics.json") + if metrics.get("status") != "success": + raise RuntimeError(f"LongBench run {path} failed: {metrics!r}.") + samples = _read_jsonl(path / "per_sample_results.jsonl") + if not samples or any(sample.get("status") != "success" for sample in samples): + raise RuntimeError(f"LongBench run {path} contains non-success samples.") + graph = _read_json(path / "decode_cuda_graph_status_rank0.json") + if not graph.get("configured_on_all_workers") or not graph.get( + "active_on_all_workers" + ): + raise RuntimeError(f"LongBench run {path} did not activate Graph everywhere.") + workers = graph.get("workers") + if not isinstance(workers, list) or not workers: + raise RuntimeError(f"LongBench run {path} has no worker status records.") + if any( + worker.get("moe_expert_provider") != "triton" + or worker.get("moe_router_provider") != "triton" + for worker in workers + ): + raise RuntimeError(f"LongBench run {path} did not use Triton providers.") + config = _read_json(path / "resolved_config.json") + expected_per_task = int(config["selection"]["samples_per_task"]) + counts = { + str(dataset): sum( + sample.get("dataset") == dataset for sample in samples + ) + for dataset in config["datasets"] + } + if any(count != expected_per_task for count in counts.values()): + raise RuntimeError( + f"LongBench run {path} has incomplete task samples: " + f"expected_per_task={expected_per_task}, counts={counts}." + ) + return { + "path": path, + "config": config, + "metrics": metrics, + "samples": samples, + "graph": graph, + "tokenizer_runtime": _read_json(path / "tokenizer_runtime.json"), + } + + +def _selection(run: dict[str, Any]) -> list[tuple[str, int]]: + return [ + (str(sample["dataset"]), int(sample["source_idx"])) + for sample in run["samples"] + ] + + +def _rendered_prompt_hashes(run: dict[str, Any]) -> list[str]: + hashes = [sample.get("rendered_prompt_sha256") for sample in run["samples"]] + if any(not isinstance(value, str) or not value for value in hashes): + raise RuntimeError( + f"LongBench run {run['path']} is missing rendered prompt hashes." + ) + return hashes + + +def _dataset_fingerprints(run: dict[str, Any]) -> dict[str, str]: + return { + str(item["dataset"]): str(item["sha256"]) + for item in run["config"]["dataset_files"] + } + + +def _task_scores(run: dict[str, Any]) -> dict[str, float]: + datasets = [str(dataset) for dataset in run["config"]["datasets"]] + scores = {} + for dataset in datasets: + score = run["metrics"].get(dataset) + if not isinstance(score, (int, float)): + raise TypeError( + f"LongBench run {run['path']} has no numeric score for {dataset}: {score!r}." + ) + scores[dataset] = float(score) + return scores + + +def _write_json(path: Path, value: Any) -> None: + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Compare the fixed Qwen3.6 MoE LongBench subset across topologies." + ) + parser.add_argument("--single", type=Path, required=True) + parser.add_argument("--pure-tp", type=Path, required=True) + parser.add_argument("--tp-ep", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + args = parser.parse_args() + runs = { + "single": _load_run(args.single.resolve()), + "pure_tp": _load_run(args.pure_tp.resolve()), + "tp_ep": _load_run(args.tp_ep.resolve()), + } + output_dir = args.output_dir.resolve() + output_dir.mkdir(parents=True, exist_ok=False) + + reference = runs["single"] + reference_selection = _selection(reference) + reference_prompt_hashes = _rendered_prompt_hashes(reference) + reference_fingerprints = _dataset_fingerprints(reference) + reference_scores = _task_scores(reference) + reference_mean = sum(reference_scores.values()) / len(reference_scores) + records: list[dict[str, Any]] = [] + for topology, run in runs.items(): + config_equal = all( + run["config"].get(key) == reference["config"].get(key) + for key in ( + "datasets", + "seed", + "decoding", + "selection", + "model_files", + "tokenizer_files", + "provider_env", + "prompt_config", + "maxlen_config", + ) + ) + same_selection = _selection(run) == reference_selection + same_rendered_prompts = ( + _rendered_prompt_hashes(run) == reference_prompt_hashes + ) + same_fingerprints = _dataset_fingerprints(run) == reference_fingerprints + same_tokenizer_runtime = ( + run["tokenizer_runtime"] == reference["tokenizer_runtime"] + ) + records.append( + { + "check": "fixed_inputs", + "topology": topology, + "status": ( + "success" + if config_equal + and same_selection + and same_rendered_prompts + and same_fingerprints + and same_tokenizer_runtime + else "metric_failed" + ), + "config_equal": config_equal, + "same_sample_ids": same_selection, + "same_rendered_prompts": same_rendered_prompts, + "same_dataset_fingerprints": same_fingerprints, + "same_tokenizer_runtime": same_tokenizer_runtime, + "num_samples": len(run["samples"]), + } + ) + scores = _task_scores(run) + for task, reference_score in reference_scores.items(): + score = scores[task] + drop = reference_score - score + records.append( + { + "check": "task_quality", + "topology": topology, + "task": task, + "status": ( + "success" if drop <= TASK_SCORE_MAX_DROP else "metric_failed" + ), + "score": score, + "single_reference_score": reference_score, + "score_drop": drop, + "max_allowed_drop": TASK_SCORE_MAX_DROP, + } + ) + mean_score = sum(scores.values()) / len(scores) + mean_drop = reference_mean - mean_score + records.append( + { + "check": "mean_quality", + "topology": topology, + "status": ( + "success" if mean_drop <= MEAN_SCORE_MAX_DROP else "metric_failed" + ), + "mean_score": mean_score, + "single_reference_mean_score": reference_mean, + "score_drop": mean_drop, + "max_allowed_drop": MEAN_SCORE_MAX_DROP, + } + ) + + failed = [record for record in records if record["status"] != "success"] + aggregate = { + "status": "success" if not failed else "metric_failed", + "num_checks": len(records), + "success_checks": len(records) - len(failed), + "failed_checks": len(failed), + "thresholds": { + "task_score_max_drop": TASK_SCORE_MAX_DROP, + "mean_score_max_drop": MEAN_SCORE_MAX_DROP, + }, + "scores": { + topology: _task_scores(run) for topology, run in runs.items() + }, + } + run_info = { + "created_at": datetime.now().isoformat(timespec="seconds"), + "command": " ".join(sys.argv), + "git_commit": subprocess.run( + ["git", "rev-parse", "HEAD"], text=True, capture_output=True, check=True + ).stdout.strip(), + "sources": {topology: str(run["path"]) for topology, run in runs.items()}, + } + _write_json(output_dir / "run_info.json", run_info) + _write_json(output_dir / "raw_outputs.json", run_info["sources"]) + _write_json(output_dir / "parsed_outputs.json", {"checks": records}) + with (output_dir / "per_sample_results.jsonl").open( + "w", encoding="utf-8" + ) as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + _write_json(output_dir / "aggregate_metrics.json", aggregate) + print(json.dumps(aggregate, ensure_ascii=False, indent=2)) + return 0 if not failed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/validation/qwen36_end_to_end_reference.py b/scripts/validation/qwen36_end_to_end_reference.py new file mode 100644 index 00000000..f62c2b8e --- /dev/null +++ b/scripts/validation/qwen36_end_to_end_reference.py @@ -0,0 +1,526 @@ +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import subprocess +import sys +import traceback +from datetime import datetime +from pathlib import Path +from typing import Any + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( + Qwen3_5MoeRMSNormGated, + torch_causal_conv1d_update, + torch_chunk_gated_delta_rule, + torch_recurrent_gated_delta_rule, +) + + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from benchmark.runtime_validation import collect_worker_runtime_status + +DEFAULT_PROMPTS = ( + "Sparse attention keeps the most useful context because", + "请用一句话解释专家并行的作用：", + "def fibonacci(n):\n \"\"\"Return the nth Fibonacci number.\"\"\"\n", +) + + +def _write_json(path: Path, value: Any) -> None: + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + with path.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n") + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _git_value(*args: str) -> str | None: + result = subprocess.run( + ["git", *args], + cwd=REPO_ROOT, + check=False, + capture_output=True, + text=True, + ) + value = result.stdout.strip() + return value or None + + +def _tokenize_prompts(tokenizer, prompts: tuple[str, ...]) -> list[list[int]]: + token_ids = [] + for prompt in prompts: + add_special_tokens = True + if tokenizer.bos_token is None or prompt.startswith(tokenizer.bos_token): + add_special_tokens = False + encoded = tokenizer.encode(prompt, add_special_tokens=add_special_tokens) + if not encoded: + raise ValueError(f"Prompt tokenized to an empty sequence: {prompt!r}.") + token_ids.append([int(token_id) for token_id in encoded]) + return token_ids + + +def _run_transformers( + args, + tokenizer, + prompt_token_ids: list[list[int]], +) -> tuple[ + list[dict[str, Any]], + list[torch.Tensor], + dict[str, Any], + list[dict[int, torch.Tensor]], + list[dict[int, torch.Tensor]], + list[dict[int, dict[str, torch.Tensor]]], +]: + model = AutoModelForCausalLM.from_pretrained( + args.model, + dtype=torch.bfloat16, + local_files_only=True, + attn_implementation=("eager" if args.torch_reference_kernels else "sdpa"), + ).to("cuda").eval() + language_model = model.model + if hasattr(language_model, "language_model"): + language_model = language_model.language_model + if args.torch_reference_kernels: + for layer in language_model.layers: + linear_attn = getattr(layer, "linear_attn", None) + if linear_attn is None: + continue + linear_attn.causal_conv1d_fn = None + linear_attn.causal_conv1d_update = torch_causal_conv1d_update + linear_attn.chunk_gated_delta_rule = torch_chunk_gated_delta_rule + linear_attn.recurrent_gated_delta_rule = ( + torch_recurrent_gated_delta_rule + ) + torch_norm = Qwen3_5MoeRMSNormGated( + linear_attn.head_v_dim, + eps=linear_attn.layer_norm_epsilon, + ).to(device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + torch_norm.weight.copy_(linear_attn.norm.weight) + linear_attn.norm = torch_norm + rows = [] + logits = [] + hidden_snapshots: list[dict[int, torch.Tensor]] = [] + cached_hidden_snapshots: list[dict[int, torch.Tensor]] = [] + selected_layers = tuple(args.debug_hidden_layers) + live_hidden: dict[int, torch.Tensor] = {} + handles = [] + + def capture(layer_idx: int): + def hook(_module, _inputs, output): + tensor = output[0] if isinstance(output, tuple) else output + live_hidden[layer_idx] = tensor[:, -1].detach().cpu() + + return hook + + if selected_layers: + for layer_idx in selected_layers: + handles.append( + language_model.layers[layer_idx].register_forward_hook( + capture(layer_idx) + ) + ) + handles.append( + language_model.norm.register_forward_hook( + capture(len(language_model.layers)) + ) + ) + try: + for sample_id, (prompt, input_ids) in enumerate( + zip(DEFAULT_PROMPTS, prompt_token_ids) + ): + input_tensor = torch.tensor( + [input_ids], dtype=torch.long, device="cuda" + ) + with torch.inference_mode(): + generated = model.generate( + input_tensor, + do_sample=False, + max_new_tokens=args.max_new_tokens, + min_new_tokens=args.max_new_tokens, + use_cache=True, + pad_token_id=tokenizer.eos_token_id, + ) + output_ids = generated[0, input_tensor.shape[1] :].tolist() + if len(output_ids) != args.max_new_tokens: + raise RuntimeError( + f"Transformers sample {sample_id} generated {len(output_ids)} " + f"tokens, expected {args.max_new_tokens}." + ) + if selected_layers: + live_hidden[-1] = ( + language_model.embed_tokens(generated[:, -2]) + .detach() + .cpu() + ) + cached_hidden_snapshots.append( + dict(sorted(live_hidden.items())) + ) + final_input = generated[:, :-1] + live_hidden.clear() + final_logits = model( + input_ids=final_input, + use_cache=False, + return_dict=True, + ).logits[0, -1].detach().cpu() + if selected_layers: + live_hidden[-1] = ( + language_model.embed_tokens(final_input[:, -1]) + .detach() + .cpu() + ) + hidden_snapshots.append(dict(sorted(live_hidden.items()))) + logits.append(final_logits) + rows.append( + { + "sample_id": sample_id, + "status": "success", + "prompt": prompt, + "prompt_token_ids": input_ids, + "output_token_ids": [int(token_id) for token_id in output_ids], + "output_text": tokenizer.decode( + output_ids, skip_special_tokens=True + ), + } + ) + finally: + for handle in handles: + handle.remove() + return ( + rows, + logits, + {"backend": "transformers", "worker_runtime_status": []}, + hidden_snapshots, + cached_hidden_snapshots, + [], + ) + + +def _run_sparsevllm( + args, + tokenizer, + prompt_token_ids: list[list[int]], +) -> tuple[ + list[dict[str, Any]], + list[torch.Tensor], + dict[str, Any], + list[dict[int, torch.Tensor]], + list[dict[int, torch.Tensor]], + list[dict[int, dict[str, torch.Tensor]]], +]: + os.environ["SPARSEVLLM_DEBUG_RUNTIME"] = "1" + os.environ["SPARSEVLLM_DEBUG_MOE"] = "1" + if args.debug_hidden_layers: + os.environ["SPARSEVLLM_DEBUG_HIDDEN_LAYERS"] = ",".join( + str(layer_idx) for layer_idx in args.debug_hidden_layers + ) + if str(REPO_ROOT / "src") not in sys.path: + sys.path.insert(0, str(REPO_ROOT / "src")) + from sparsevllm import LLM, SamplingParams + + max_prompt_len = max(len(item) for item in prompt_token_ids) + llm = LLM( + model=str(args.model), + tensor_parallel_size=args.tensor_parallel_size, + expert_parallel_size=args.expert_parallel_size, + data_parallel_size=1, + enforce_eager=not args.decode_cuda_graph, + decode_cuda_graph=args.decode_cuda_graph, + gpu_memory_utilization=args.gpu_memory_utilization, + weight_loading_workers=args.weight_loading_workers, + max_model_len=max_prompt_len + args.max_new_tokens + 32, + max_num_seqs_in_batch=1, + max_decoding_seqs=1, + engine_prefill_chunk_size=max(64, max_prompt_len), + enable_profiler=False, + ) + rows = [] + logits = [] + debug_summaries = [] + hidden_snapshots = [] + moe_snapshots = [] + try: + sampling_params = SamplingParams( + temperature=0.0, + top_p=1.0, + max_tokens=args.max_new_tokens, + ignore_eos=True, + ) + for sample_id, (prompt, input_ids) in enumerate( + zip(DEFAULT_PROMPTS, prompt_token_ids) + ): + result = llm.generate( + [input_ids], sampling_params, use_tqdm=False + )[0] + output_ids = [int(token_id) for token_id in result["token_ids"]] + if len(output_ids) != args.max_new_tokens: + raise RuntimeError( + f"Sparse-vLLM sample {sample_id} generated {len(output_ids)} " + f"tokens, expected {args.max_new_tokens}." + ) + logits.append(llm.debug_last_logits().detach().cpu()[0]) + debug_summaries.append(llm.debug_sparse_state_summaries()) + if args.debug_hidden_layers: + hidden_snapshots.append(llm.debug_hidden_states()) + if args.debug_moe_states: + moe_snapshots.append(llm.debug_moe_states()) + rows.append( + { + "sample_id": sample_id, + "status": "success", + "prompt": prompt, + "prompt_token_ids": input_ids, + "output_token_ids": output_ids, + "output_text": result["text"], + } + ) + worker_status = collect_worker_runtime_status(llm) + if args.decode_cuda_graph and not all( + bool(status["decode_cuda_graph_active"]) for status in worker_status + ): + raise RuntimeError( + "Decode CUDA Graph is not active on every rank: " + f"{worker_status!r}." + ) + finally: + llm.exit() + return ( + rows, + logits, + { + "backend": "sparsevllm", + "worker_runtime_status": worker_status, + "debug_sparse_state_summaries": debug_summaries, + }, + hidden_snapshots, + hidden_snapshots, + moe_snapshots, + ) + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Produce deterministic Qwen3.6 MoE end-to-end reference artifacts." + ) + parser.add_argument("--backend", choices=("transformers", "sparsevllm"), required=True) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--max-new-tokens", type=int, default=8) + parser.add_argument("--seed", type=int, default=20260810) + parser.add_argument("--tensor-parallel-size", type=int, default=1) + parser.add_argument("--expert-parallel-size", type=int, default=1) + parser.add_argument("--decode-cuda-graph", action="store_true") + parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) + parser.add_argument("--weight-loading-workers", type=int, default=16) + parser.add_argument( + "--torch-reference-kernels", + action="store_true", + help=( + "For the Transformers backend, force eager full attention and the " + "explicit Torch Gated DeltaNet conv/chunk/recurrent/norm functions." + ), + ) + parser.add_argument( + "--debug-hidden-layers", + type=int, + nargs="*", + default=(), + help="Capture last-token hidden states after the selected decoder layers.", + ) + parser.add_argument( + "--debug-moe-states", + action="store_true", + help="Capture rank-0 per-layer MoE inputs, routing, and outputs.", + ) + parser.add_argument( + "--forced-prefix-artifact", + type=Path, + default=None, + help=( + "Append all but the last generated token from a prior " + "per_sample_results.jsonl artifact to each fixed prompt." + ), + ) + args = parser.parse_args() + if args.max_new_tokens <= 0: + raise ValueError("--max-new-tokens must be positive.") + if args.torch_reference_kernels and args.backend != "transformers": + raise ValueError("--torch-reference-kernels requires --backend transformers.") + invalid_hidden_layers = [ + layer_idx + for layer_idx in args.debug_hidden_layers + if layer_idx < 0 or layer_idx >= 40 + ] + if invalid_hidden_layers: + raise ValueError( + "--debug-hidden-layers must be in [0, 39], got " + f"{invalid_hidden_layers}." + ) + if not torch.cuda.is_available(): + raise RuntimeError("This validation requires CUDA.") + + args.model = args.model.resolve() + args.output_dir = args.output_dir.resolve() + if args.forced_prefix_artifact is not None: + args.forced_prefix_artifact = args.forced_prefix_artifact.resolve() + if not args.forced_prefix_artifact.is_file(): + raise FileNotFoundError( + "Forced-prefix artifact does not exist: " + f"{args.forced_prefix_artifact}." + ) + args.output_dir.mkdir(parents=True, exist_ok=False) + run_info = { + "created_at": datetime.now().isoformat(timespec="seconds"), + "command": " ".join(sys.argv), + "backend": args.backend, + "model": str(args.model), + "model_config_sha256": _sha256(args.model / "config.json"), + "model_index_sha256": _sha256( + args.model / "model.safetensors.index.json" + ), + "seed": args.seed, + "max_new_tokens": args.max_new_tokens, + "temperature": 0.0, + "top_p": 1.0, + "tensor_parallel_size": args.tensor_parallel_size, + "expert_parallel_size": args.expert_parallel_size, + "decode_cuda_graph": args.decode_cuda_graph, + "torch_reference_kernels": args.torch_reference_kernels, + "forced_prefix_artifact": ( + None + if args.forced_prefix_artifact is None + else str(args.forced_prefix_artifact) + ), + "forced_prefix_artifact_sha256": ( + None + if args.forced_prefix_artifact is None + else _sha256(args.forced_prefix_artifact) + ), + "gpu": torch.cuda.get_device_name(0), + "git_commit": _git_value("rev-parse", "HEAD"), + "git_branch": _git_value("branch", "--show-current"), + "git_dirty": bool(_git_value("status", "--porcelain")), + "requested_moe_provider": os.getenv("SPARSEVLLM_MOE_PROVIDER", "auto"), + "requested_moe_router_provider": os.getenv( + "SPARSEVLLM_MOE_ROUTER_PROVIDER", "auto" + ), + } + _write_json(args.output_dir / "run_info.json", run_info) + tokenizer = AutoTokenizer.from_pretrained(args.model, local_files_only=True) + prompt_token_ids = _tokenize_prompts(tokenizer, DEFAULT_PROMPTS) + if args.forced_prefix_artifact is not None: + forced_rows = [ + json.loads(line) + for line in args.forced_prefix_artifact.read_text( + encoding="utf-8" + ).splitlines() + if line.strip() + ] + if len(forced_rows) != len(prompt_token_ids): + raise ValueError( + "Forced-prefix artifact must contain one row per fixed prompt: " + f"expected={len(prompt_token_ids)} got={len(forced_rows)}." + ) + for sample_id, (token_ids, row) in enumerate( + zip(prompt_token_ids, forced_rows) + ): + output_token_ids = row.get("output_token_ids") + if not isinstance(output_token_ids, list) or len(output_token_ids) < 2: + raise ValueError( + f"Forced-prefix sample {sample_id} needs at least two " + "output_token_ids." + ) + token_ids.extend(int(token_id) for token_id in output_token_ids[:-1]) + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + + try: + if args.backend == "transformers": + ( + rows, + logits, + parsed, + hidden_snapshots, + cached_hidden_snapshots, + moe_snapshots, + ) = _run_transformers( + args, tokenizer, prompt_token_ids + ) + else: + ( + rows, + logits, + parsed, + hidden_snapshots, + cached_hidden_snapshots, + moe_snapshots, + ) = _run_sparsevllm( + args, tokenizer, prompt_token_ids + ) + except Exception as exc: + failure = { + "sample_id": None, + "status": "model_failed", + "error": repr(exc), + "traceback": traceback.format_exc(), + } + _write_jsonl(args.output_dir / "raw_outputs.jsonl", [failure]) + _write_jsonl(args.output_dir / "parsed_outputs.jsonl", [failure]) + _write_jsonl(args.output_dir / "per_sample_results.jsonl", [failure]) + _write_json(args.output_dir / "aggregate_metrics.json", failure) + raise + + _write_jsonl(args.output_dir / "raw_outputs.jsonl", rows) + _write_jsonl(args.output_dir / "parsed_outputs.jsonl", rows) + _write_jsonl(args.output_dir / "per_sample_results.jsonl", rows) + torch.save(logits, args.output_dir / "raw_logits.pt") + if hidden_snapshots: + torch.save( + hidden_snapshots, + args.output_dir / "raw_hidden_states.pt", + ) + if cached_hidden_snapshots: + torch.save( + cached_hidden_snapshots, + args.output_dir / "raw_cached_hidden_states.pt", + ) + if moe_snapshots: + torch.save( + moe_snapshots, + args.output_dir / "raw_moe_states.pt", + ) + _write_json(args.output_dir / "runtime_status.json", parsed) + aggregate = { + "status": "success", + "num_samples": len(rows), + "success_samples": sum(row["status"] == "success" for row in rows), + "failed_samples": sum(row["status"] != "success" for row in rows), + } + _write_json(args.output_dir / "aggregate_metrics.json", aggregate) + print(json.dumps(aggregate, ensure_ascii=False, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/validation/qwen36_moe_bf16_reference.py b/scripts/validation/qwen36_moe_bf16_reference.py new file mode 100644 index 00000000..a4db852d --- /dev/null +++ b/scripts/validation/qwen36_moe_bf16_reference.py @@ -0,0 +1,367 @@ +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import subprocess +import sys +from datetime import datetime +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F +from safetensors import safe_open +from transformers import AutoConfig +from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( + Qwen3_5MoeSparseMoeBlock, +) + + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT / "src") not in sys.path: + sys.path.insert(0, str(REPO_ROOT / "src")) + +from sparsevllm.operators.moe import MoeOpSpec, resolve_moe_provider +from sparsevllm.operators.moe_router import ( + MoeRouterOpSpec, + resolve_moe_router_provider, +) + + +WEIGHT_PREFIX = "model.language_model.layers.0.mlp." +WEIGHT_NAMES = ( + "gate.weight", + "experts.gate_up_proj", + "experts.down_proj", + "shared_expert.gate_proj.weight", + "shared_expert.up_proj.weight", + "shared_expert.down_proj.weight", + "shared_expert_gate.weight", +) + + +def _write_json(path: Path, value: Any) -> None: + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _git_value(*args: str) -> str | None: + result = subprocess.run( + ["git", *args], + cwd=REPO_ROOT, + check=False, + capture_output=True, + text=True, + ) + value = result.stdout.strip() + return value or None + + +def _load_layer_weights(model_path: Path) -> dict[str, torch.Tensor]: + index_path = model_path / "model.safetensors.index.json" + if not index_path.is_file(): + raise FileNotFoundError(f"Missing checkpoint index: {index_path}") + weight_map = json.loads(index_path.read_text(encoding="utf-8"))["weight_map"] + loaded: dict[str, torch.Tensor] = {} + by_shard: dict[str, list[tuple[str, str]]] = {} + for local_name in WEIGHT_NAMES: + checkpoint_name = WEIGHT_PREFIX + local_name + shard = weight_map.get(checkpoint_name) + if shard is None: + raise KeyError(f"Missing checkpoint tensor {checkpoint_name!r}.") + by_shard.setdefault(shard, []).append((checkpoint_name, local_name)) + for shard, names in by_shard.items(): + shard_path = model_path / shard + if not shard_path.is_file(): + raise FileNotFoundError(f"Missing checkpoint shard: {shard_path}") + with safe_open(shard_path, framework="pt", device="cpu") as handle: + for checkpoint_name, local_name in names: + tensor = handle.get_tensor(checkpoint_name) + if tensor.dtype != torch.bfloat16: + raise TypeError( + f"{checkpoint_name} must be BF16, got {tensor.dtype}." + ) + loaded[local_name] = tensor + return loaded + + +def _tensor_summary(value: torch.Tensor) -> dict[str, Any]: + finite = torch.isfinite(value) + float_value = value.float() + return { + "shape": list(value.shape), + "dtype": str(value.dtype), + "finite": bool(finite.all().item()), + "min": float(float_value.min().item()), + "max": float(float_value.max().item()), + "mean": float(float_value.mean().item()), + } + + +def _error_metrics( + actual: torch.Tensor, + expected: torch.Tensor, +) -> dict[str, float]: + difference = (actual.float() - expected.float()).abs() + denominator = expected.float().abs().clamp_min(1.0e-6) + return { + "max_abs": float(difference.max().item()), + "mean_abs": float(difference.mean().item()), + "max_rel": float((difference / denominator).max().item()), + } + + +def _record_close( + name: str, + actual: torch.Tensor, + expected: torch.Tensor, + *, + atol: float, + rtol: float, +) -> dict[str, Any]: + metrics = _error_metrics(actual, expected) + passed = bool(torch.allclose(actual.float(), expected.float(), atol=atol, rtol=rtol)) + return { + "component": name, + "status": "success" if passed else "metric_failed", + "atol": float(atol), + "rtol": float(rtol), + **metrics, + } + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Compare Qwen3.6 MoE BF16 checkpoint math with Transformers." + ) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--tokens", type=int, default=17) + parser.add_argument("--seed", type=int, default=20260810) + args = parser.parse_args() + if args.tokens <= 0: + raise ValueError(f"--tokens must be positive, got {args.tokens}.") + if not torch.cuda.is_available(): + raise RuntimeError("This validation requires one CUDA device.") + + model_path = args.model.resolve() + output_dir = args.output_dir.resolve() + output_dir.mkdir(parents=True, exist_ok=False) + index_path = model_path / "model.safetensors.index.json" + config_path = model_path / "config.json" + run_info = { + "created_at": datetime.now().isoformat(timespec="seconds"), + "command": " ".join(sys.argv), + "model": str(model_path), + "model_config_sha256": _sha256(config_path), + "model_index_sha256": _sha256(index_path), + "seed": int(args.seed), + "tokens": int(args.tokens), + "dtype": "torch.bfloat16", + "device": torch.cuda.get_device_name(0), + "git_commit": _git_value("rev-parse", "HEAD"), + "git_branch": _git_value("branch", "--show-current"), + "git_dirty": bool(_git_value("status", "--porcelain")), + "tolerances": { + "routing_weights": {"atol": 0.004, "rtol": 0.004}, + "routed_experts": {"atol": 0.05, "rtol": 0.05}, + "shared_expert": {"atol": 0.02, "rtol": 0.02}, + "moe_output": {"atol": 0.05, "rtol": 0.05}, + }, + "requested_moe_provider": os.getenv("SPARSEVLLM_MOE_PROVIDER", "auto"), + "requested_moe_router_provider": os.getenv( + "SPARSEVLLM_MOE_ROUTER_PROVIDER", "auto" + ), + } + + weights = _load_layer_weights(model_path) + outer_config = AutoConfig.from_pretrained(model_path, local_files_only=True) + config = outer_config.text_config + if config.torch_dtype != torch.bfloat16: + raise TypeError(f"Reference config must be BF16, got {config.torch_dtype}.") + num_experts = int(config.num_experts) + expert_intermediate_size = int(weights["experts.gate_up_proj"].shape[1] // 2) + moe_spec = MoeOpSpec( + num_experts=num_experts, + num_local_experts=num_experts, + hidden_size=int(config.hidden_size), + intermediate_size=expert_intermediate_size, + top_k=int(config.num_experts_per_tok), + activation_dtype=torch.bfloat16, + weight_dtype=torch.bfloat16, + block_shape=None, + ep_size=1, + cuda_graph=False, + tp_size=1, + routing_method="softmax", + ) + moe_provider = resolve_moe_provider(moe_spec, device_index=0) + router_spec = MoeRouterOpSpec( + num_experts=num_experts, + top_k=int(config.num_experts_per_tok), + activation_dtype=torch.bfloat16, + norm_topk_prob=True, + cuda_graph=False, + ) + router_provider = resolve_moe_router_provider(router_spec, device_index=0) + run_info["resolved_moe_provider"] = moe_provider.name + run_info["resolved_moe_router_provider"] = router_provider.name + _write_json(output_dir / "run_info.json", run_info) + + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + reference = Qwen3_5MoeSparseMoeBlock(config) + finally: + torch.set_default_dtype(previous_dtype) + missing, unexpected = reference.load_state_dict(weights, strict=False) + if missing or unexpected: + raise RuntimeError( + f"Transformers MoE weight mismatch: missing={missing}, unexpected={unexpected}." + ) + reference = reference.to(device="cuda", dtype=torch.bfloat16).eval() + weights = { + name: tensor.to(device="cuda", non_blocking=False) + for name, tensor in weights.items() + } + + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + hidden = torch.randn( + (args.tokens, int(config.hidden_size)), + device="cuda", + dtype=torch.bfloat16, + ) + with torch.inference_mode(): + reference_logits, reference_routing_weights, reference_ids = reference.gate( + hidden + ) + reference_routed = reference.experts( + hidden, reference_ids, reference_routing_weights + ) + reference_shared = reference.shared_expert(hidden) + reference_shared = torch.sigmoid(reference.shared_expert_gate(hidden)) * reference_shared + reference_output = reference_routed + reference_shared + + actual_logits = F.linear(hidden, weights["gate.weight"]) + actual_routing_weights, actual_ids = router_provider.run( + router_spec, + actual_logits, + ) + actual_routed = moe_provider.run( + moe_spec, + hidden, + actual_ids, + actual_routing_weights, + weights["experts.gate_up_proj"], + weights["experts.down_proj"], + None, + None, + local_expert_start=0, + ep_rank=0, + ) + shared_gate = F.linear(hidden, weights["shared_expert.gate_proj.weight"]) + shared_up = F.linear(hidden, weights["shared_expert.up_proj.weight"]) + actual_shared = F.linear( + F.silu(shared_gate) * shared_up, + weights["shared_expert.down_proj.weight"], + ) + actual_shared *= torch.sigmoid( + F.linear(hidden, weights["shared_expert_gate.weight"]) + ) + actual_output = actual_routed + actual_shared + torch.cuda.synchronize() + + records = [ + _record_close( + "router_logits", + actual_logits, + reference_logits, + atol=0.0, + rtol=0.0, + ), + { + "component": "routing_ids", + "status": "success" if torch.equal(actual_ids, reference_ids.to(torch.int32)) else "metric_failed", + "mismatch_count": int((actual_ids != reference_ids.to(torch.int32)).sum().item()), + }, + _record_close( + "routing_weights", + actual_routing_weights, + reference_routing_weights, + atol=0.004, + rtol=0.004, + ), + _record_close( + "routed_experts", + actual_routed, + reference_routed, + atol=0.05, + rtol=0.05, + ), + _record_close( + "shared_expert", + actual_shared, + reference_shared, + atol=0.02, + rtol=0.02, + ), + _record_close( + "moe_output", + actual_output, + reference_output, + atol=0.05, + rtol=0.05, + ), + ] + raw_outputs = { + "hidden": hidden.cpu(), + "reference_logits": reference_logits.cpu(), + "actual_logits": actual_logits.cpu(), + "reference_ids": reference_ids.cpu(), + "actual_ids": actual_ids.cpu(), + "reference_routing_weights": reference_routing_weights.cpu(), + "actual_routing_weights": actual_routing_weights.cpu(), + "reference_routed": reference_routed.cpu(), + "actual_routed": actual_routed.cpu(), + "reference_shared": reference_shared.cpu(), + "actual_shared": actual_shared.cpu(), + "reference_output": reference_output.cpu(), + "actual_output": actual_output.cpu(), + } + torch.save(raw_outputs, output_dir / "raw_outputs.pt") + parsed = { + name: _tensor_summary(tensor) + for name, tensor in raw_outputs.items() + } + _write_json(output_dir / "parsed_outputs.json", parsed) + with (output_dir / "per_sample_results.jsonl").open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n") + failed = [record for record in records if record["status"] != "success"] + aggregate = { + "status": "success" if not failed else "metric_failed", + "num_checks": len(records), + "success_checks": len(records) - len(failed), + "failed_checks": len(failed), + "records": records, + } + _write_json(output_dir / "aggregate_metrics.json", aggregate) + print(json.dumps(aggregate, ensure_ascii=False, indent=2)) + return 0 if not failed else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/validation/qwen36_summarize_microbench.py b/scripts/validation/qwen36_summarize_microbench.py new file mode 100644 index 00000000..18b16b2c --- /dev/null +++ b/scripts/validation/qwen36_summarize_microbench.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +import argparse +import json +import statistics +import subprocess +import sys +from datetime import datetime +from pathlib import Path +from typing import Any + + +CASES = ( + (1024, 1, 512), + (4096, 8, 512), + (32768, 2, 512), + (65536, 1, 1024), + (131072, 1, 1024), +) +TOPOLOGIES = ("single", "pure_tp", "tp_ep") +MODES = ("graph", "eager") +METRICS = ( + "ttft_s", + "prefill_tok_s", + "decode_tok_s", + "itl_ms", + "end_to_end_tok_s", + "peak_memory_gb", +) + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + rows = [] + for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + value = json.loads(line) + if not isinstance(value, dict): + raise TypeError(f"Expected an object at {path}:{line_no}.") + rows.append(value) + return rows + + +def _write_json(path: Path, value: Any) -> None: + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _stats(rows: list[dict[str, Any]], metric: str) -> dict[str, float]: + values = [float(row[metric]) for row in rows] + return { + "median": statistics.median(values), + "min": min(values), + "max": max(values), + "relative_range": ( + (max(values) - min(values)) / statistics.median(values) + if statistics.median(values) + else 0.0 + ), + } + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Validate and summarize a completed Qwen3.6 microbench matrix." + ) + parser.add_argument("--matrix-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=2) + parser.add_argument( + "--modes", + default="graph,eager", + help="Comma-separated subset of graph,eager.", + ) + args = parser.parse_args() + modes = tuple(part.strip() for part in args.modes.split(",") if part.strip()) + if not modes or len(set(modes)) != len(modes) or any( + mode not in MODES for mode in modes + ): + raise ValueError( + "--modes must contain each of 'graph' and 'eager' at most once, " + f"got {args.modes!r}." + ) + matrix_dir = args.matrix_dir.resolve() + output_dir = args.output_dir.resolve() + output_dir.mkdir(parents=True, exist_ok=False) + records = _read_jsonl(matrix_dir / "per_sample_results.jsonl") + + expected_keys = { + (topology, mode, length, batch, output) + for length, batch, output in CASES + for topology in TOPOLOGIES + for mode in modes + } + groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {} + validation_records = [] + for record in records: + key = ( + record.get("topology"), + record.get("mode"), + int(record.get("prompt_tokens")), + int(record.get("batch_size")), + int(record.get("max_new_tokens")), + ) + groups.setdefault(key, []).append(record) + workers = record.get("worker_runtime_status") or [] + expected_workers = 1 if record.get("topology") == "single" else 2 + graph_active = ( + record.get("mode") != "graph" + or ( + len(workers) == expected_workers + and all(worker.get("decode_cuda_graph_active") for worker in workers) + ) + ) + providers_ok = ( + record.get("moe_expert_provider") == "triton" + and record.get("moe_router_provider") == "triton" + and all( + worker.get("moe_expert_provider") == "triton" + and worker.get("moe_router_provider") == "triton" + for worker in workers + ) + ) + validation_records.append( + { + "check": "measured_run", + "run_name": record.get("run_name"), + "status": ( + "success" + if record.get("status") == "success" + and graph_active + and providers_ok + else "metric_failed" + ), + "graph_active_on_all_workers": graph_active, + "providers_ok": providers_ok, + "worker_count": len(workers), + } + ) + + missing_groups = sorted(expected_keys - set(groups)) + unexpected_groups = sorted(set(groups) - expected_keys) + summaries = [] + for key in sorted(expected_keys): + rows = groups.get(key, []) + successful = [row for row in rows if row.get("status") == "success"] + summaries.append( + { + "topology": key[0], + "mode": key[1], + "prompt_tokens": key[2], + "batch_size": key[3], + "max_new_tokens": key[4], + "status": ( + "success" + if len(successful) == args.repeats + else "metric_failed" + ), + "successful_repeats": len(successful), + "required_repeats": args.repeats, + "metrics": ( + {metric: _stats(successful, metric) for metric in METRICS} + if len(successful) == args.repeats + else None + ), + } + ) + + summary_by_key = { + ( + row["topology"], + row["mode"], + row["prompt_tokens"], + row["batch_size"], + row["max_new_tokens"], + ): row + for row in summaries + } + speedups = [] + if {"graph", "eager"} <= set(modes): + for length, batch, output in CASES: + for topology in TOPOLOGIES: + graph = summary_by_key[(topology, "graph", length, batch, output)] + eager = summary_by_key[(topology, "eager", length, batch, output)] + if graph["status"] != "success" or eager["status"] != "success": + speedups.append( + { + "topology": topology, + "prompt_tokens": length, + "batch_size": batch, + "max_new_tokens": output, + "status": "metric_failed", + } + ) + continue + graph_decode = graph["metrics"]["decode_tok_s"]["median"] + eager_decode = eager["metrics"]["decode_tok_s"]["median"] + graph_itl = graph["metrics"]["itl_ms"]["median"] + eager_itl = eager["metrics"]["itl_ms"]["median"] + speedups.append( + { + "topology": topology, + "prompt_tokens": length, + "batch_size": batch, + "max_new_tokens": output, + "status": "success", + "decode_throughput_speedup": graph_decode / eager_decode, + "itl_speedup": eager_itl / graph_itl, + "graph_decode_tok_s": graph_decode, + "eager_decode_tok_s": eager_decode, + } + ) + + failed = [ + record + for record in [*validation_records, *summaries, *speedups] + if record["status"] != "success" + ] + aggregate = { + "status": ( + "success" + if not failed and not missing_groups and not unexpected_groups + else "metric_failed" + ), + "num_runs": len(records), + "expected_runs": len(expected_keys) * args.repeats, + "missing_groups": missing_groups, + "unexpected_groups": unexpected_groups, + "failed_checks": len(failed), + "summaries": summaries, + "graph_speedups": speedups, + } + run_info = { + "created_at": datetime.now().isoformat(timespec="seconds"), + "command": " ".join(sys.argv), + "git_commit": subprocess.run( + ["git", "rev-parse", "HEAD"], text=True, capture_output=True, check=True + ).stdout.strip(), + "matrix_dir": str(matrix_dir), + "repeats": args.repeats, + "modes": list(modes), + } + _write_json(output_dir / "run_info.json", run_info) + _write_json(output_dir / "raw_outputs.json", {"matrix_dir": str(matrix_dir)}) + _write_json( + output_dir / "parsed_outputs.json", + {"run_checks": validation_records, "summaries": summaries, "speedups": speedups}, + ) + with (output_dir / "per_sample_results.jsonl").open( + "w", encoding="utf-8" + ) as handle: + for record in [*validation_records, *summaries, *speedups]: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + _write_json(output_dir / "aggregate_metrics.json", aggregate) + print(json.dumps(aggregate, ensure_ascii=False, indent=2)) + return 0 if aggregate["status"] == "success" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/validation/run_qwen36_microbench_matrix.py b/scripts/validation/run_qwen36_microbench_matrix.py new file mode 100644 index 00000000..79ce3dd3 --- /dev/null +++ b/scripts/validation/run_qwen36_microbench_matrix.py @@ -0,0 +1,391 @@ +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import statistics +import subprocess +import sys +import time +from datetime import datetime +from pathlib import Path +from typing import Any + + +CASES = ( + (1024, 1, 512), + (4096, 8, 512), + (32768, 2, 512), + (65536, 1, 1024), + (131072, 1, 1024), +) +TOPOLOGIES = ( + ("single", "0", 1, 1), + ("pure_tp", "0,1", 2, 1), + ("tp_ep", "0,1", 2, 2), +) +MODES = ("graph", "eager") +METRICS = ( + "ttft_s", + "prefill_tok_s", + "decode_tok_s", + "itl_ms", + "end_to_end_tok_s", + "peak_memory_gb", +) + + +def _write_json(path: Path, value: Any) -> None: + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _gpu_snapshot() -> dict[str, Any]: + processes = subprocess.run( + [ + "nvidia-smi", + "--query-compute-apps=gpu_uuid,pid,process_name,used_memory", + "--format=csv,noheader,nounits", + ], + text=True, + capture_output=True, + check=True, + ).stdout.strip() + gpu_lines = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu", + "--format=csv,noheader,nounits", + ], + text=True, + capture_output=True, + check=True, + ).stdout.strip().splitlines() + gpus = [] + for line in gpu_lines: + fields = [part.strip() for part in line.split(",")] + if len(fields) != 6: + raise RuntimeError(f"Unexpected nvidia-smi GPU row: {line!r}.") + gpus.append( + { + "index": int(fields[0]), + "uuid": fields[1], + "name": fields[2], + "memory_used_mib": int(fields[3]), + "memory_total_mib": int(fields[4]), + "utilization_percent": int(fields[5]), + } + ) + return {"compute_processes": processes.splitlines() if processes else [], "gpus": gpus} + + +def _wait_for_all_devices_idle(timeout_s: int) -> dict[str, Any]: + deadline = time.monotonic() + timeout_s + while True: + snapshot = _gpu_snapshot() + idle = not snapshot["compute_processes"] and all( + gpu["utilization_percent"] <= 1 for gpu in snapshot["gpus"] + ) + if idle: + return snapshot + if time.monotonic() >= deadline: + raise TimeoutError( + f"GPUs remained busy for {timeout_s}s; last snapshot={snapshot!r}." + ) + print(f"[gpu-idle] busy; waiting 10s: {snapshot}", flush=True) + time.sleep(10) + + +def _read_single_record(run_dir: Path) -> dict[str, Any]: + aggregate_path = run_dir / "aggregate_metrics.json" + if not aggregate_path.is_file(): + return { + "status": "model_failed", + "error": f"missing artifact {aggregate_path}", + } + aggregate = json.loads(aggregate_path.read_text(encoding="utf-8")) + records = aggregate.get("records") + if not isinstance(records, list) or len(records) != 1: + return { + "status": "model_failed", + "error": f"expected exactly one microbench record, got {records!r}", + } + return records[0] + + +def _aggregate(records: list[dict[str, Any]], repeats: int) -> dict[str, Any]: + groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {} + for record in records: + key = ( + record["topology"], + record["mode"], + record["prompt_tokens"], + record["batch_size"], + record["max_new_tokens"], + ) + groups.setdefault(key, []).append(record) + summaries = [] + failed_groups = 0 + for key, rows in sorted(groups.items()): + successful = [row for row in rows if row.get("status") == "success"] + status = "success" if len(successful) == repeats else "model_failed" + if status != "success": + failed_groups += 1 + metrics = {} + for metric in METRICS: + values = [float(row[metric]) for row in successful if metric in row] + metrics[metric] = ( + None + if len(values) != repeats + else { + "median": statistics.median(values), + "min": min(values), + "max": max(values), + } + ) + summaries.append( + { + "topology": key[0], + "mode": key[1], + "prompt_tokens": key[2], + "batch_size": key[3], + "max_new_tokens": key[4], + "status": status, + "successful_repeats": len(successful), + "required_repeats": repeats, + "metrics": metrics, + } + ) + return { + "status": "success" if not failed_groups else "model_failed", + "num_groups": len(summaries), + "successful_groups": len(summaries) - failed_groups, + "failed_groups": failed_groups, + "required_repeats": repeats, + "summaries": summaries, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description="Run the fixed Qwen3.6 MoE benchmark matrix.") + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=2) + parser.add_argument( + "--modes", + default="graph,eager", + help="Comma-separated subset of graph,eager.", + ) + parser.add_argument("--idle-timeout", type=int, default=600) + parser.add_argument("--run-timeout", type=int, default=3600) + args = parser.parse_args() + if args.repeats < 2: + raise ValueError("The acceptance matrix requires at least two measured repeats.") + if args.run_timeout <= 0: + raise ValueError("--run-timeout must be positive.") + modes = tuple(part.strip() for part in args.modes.split(",") if part.strip()) + if not modes or len(set(modes)) != len(modes) or any( + mode not in MODES for mode in modes + ): + raise ValueError( + "--modes must contain each of 'graph' and 'eager' at most once, " + f"got {args.modes!r}." + ) + model = args.model.resolve() + output_dir = args.output_dir.resolve() + output_dir.mkdir(parents=True, exist_ok=False) + for name in ("config.json", "model.safetensors.index.json"): + if not (model / name).is_file(): + raise FileNotFoundError(f"Required model metadata is missing: {model / name}.") + + initial_gpu = _wait_for_all_devices_idle(args.idle_timeout) + run_info = { + "created_at": datetime.now().isoformat(timespec="seconds"), + "command": " ".join([sys.executable, *sys.argv]), + "model": str(model), + "model_config_sha256": _sha256(model / "config.json"), + "model_index_sha256": _sha256(model / "model.safetensors.index.json"), + "git_commit": subprocess.run( + ["git", "rev-parse", "HEAD"], text=True, capture_output=True, check=True + ).stdout.strip(), + "git_branch": subprocess.run( + ["git", "branch", "--show-current"], text=True, capture_output=True, check=True + ).stdout.strip(), + "git_dirty": bool( + subprocess.run( + ["git", "status", "--porcelain"], + text=True, + capture_output=True, + check=True, + ).stdout.strip() + ), + "cases": [ + {"prompt_tokens": length, "batch_size": batch, "max_new_tokens": output} + for length, batch, output in CASES + ], + "topologies": [ + { + "name": name, + "cuda_visible_devices": devices, + "outer_tp_size": tp, + "expert_parallel_size": ep, + "moe_tp_size": tp // ep, + } + for name, devices, tp, ep in TOPOLOGIES + ], + "modes": list(modes), + "repeats": args.repeats, + "seed": 20260810, + "warmup_output_len": 8, + "run_timeout_s": args.run_timeout, + "initial_gpu_snapshot": initial_gpu, + } + _write_json(output_dir / "run_info.json", run_info) + + records: list[dict[str, Any]] = [] + total = len(CASES) * len(TOPOLOGIES) * len(modes) * args.repeats + run_index = 0 + for length, batch, max_new_tokens in CASES: + for topology, devices, tp_size, ep_size in TOPOLOGIES: + for mode in modes: + for repeat in range(1, args.repeats + 1): + run_index += 1 + idle_snapshot = _wait_for_all_devices_idle(args.idle_timeout) + run_name = ( + f"{length}_bs{batch}_out{max_new_tokens}_" + f"{topology}_{mode}_r{repeat}" + ) + run_dir = output_dir / "runs" / run_name + hyper_params = { + "tensor_parallel_size": tp_size, + "expert_parallel_size": ep_size, + "enforce_eager": mode == "eager", + "decode_cuda_graph": mode == "graph", + "gpu_memory_utilization": 0.9, + "engine_prefill_chunk_size": 8192, + "max_num_batched_tokens": 65536, + "weight_loading_workers": 16, + } + command = [ + sys.executable, + "benchmark/microbench.py", + "--model_path", + str(model), + "--lengths", + str(length), + "--batch_sizes", + str(batch), + "--output_len", + str(max_new_tokens), + "--methods", + "vanilla", + "--temperature", + "0", + "--top_p", + "1", + "--seed", + "20260810", + "--synchronize_step_timing", + "--warmup_output_len", + "8", + "--hyper_params", + json.dumps(hyper_params, separators=(",", ":")), + "--output_dir", + str(run_dir), + ] + env = os.environ.copy() + env.update( + { + "CUDA_VISIBLE_DEVICES": devices, + "PYTHONUNBUFFERED": "1", + "SPARSEVLLM_MOE_PROVIDER": "triton", + "SPARSEVLLM_MOE_ROUTER_PROVIDER": "triton", + } + ) + print( + f"[matrix {run_index}/{total}] {run_name}: {' '.join(command)}", + flush=True, + ) + timed_out = False + try: + process = subprocess.run( + command, + cwd=Path(__file__).resolve().parents[2], + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + check=False, + timeout=args.run_timeout, + ) + console_output = process.stdout + returncode = process.returncode + except subprocess.TimeoutExpired as exc: + timed_out = True + console_output = exc.stdout or "" + if isinstance(console_output, bytes): + console_output = console_output.decode( + "utf-8", errors="replace" + ) + console_output += ( + "\nMatrix runner timed out this run after " + f"{args.run_timeout}s.\n" + ) + returncode = None + run_dir.mkdir(parents=True, exist_ok=True) + (run_dir / "console.log").write_text( + console_output, encoding="utf-8" + ) + record = _read_single_record(run_dir) + record.update( + { + "run_name": run_name, + "topology": topology, + "mode": mode, + "repeat": repeat, + "prompt_tokens": length, + "batch_size": batch, + "max_new_tokens": max_new_tokens, + "returncode": returncode, + "timed_out": timed_out, + "command": command, + "idle_snapshot_before": idle_snapshot, + "artifact_dir": str(run_dir), + } + ) + if timed_out or returncode != 0: + record["status"] = "model_failed" + records.append(record) + with (output_dir / "per_sample_results.jsonl").open( + "a", encoding="utf-8" + ) as handle: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + print( + f"[matrix {run_index}/{total}] status={record['status']}", + flush=True, + ) + + aggregate = _aggregate(records, args.repeats) + _write_json(output_dir / "aggregate_metrics.json", aggregate) + _write_json(output_dir / "parsed_outputs.json", {"records": records}) + _write_json( + output_dir / "raw_outputs.json", + {"run_artifact_dirs": [record["artifact_dir"] for record in records]}, + ) + print(json.dumps(aggregate, ensure_ascii=False, indent=2), flush=True) + return 0 if aggregate["status"] == "success" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/sparsevllm/configs/model.py b/src/sparsevllm/configs/model.py index 0200316d..7c148d56 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -41,10 +41,10 @@ def _load_raw_qwen35_config(model_path: str, error: Exception) -> SimpleNamespac ) from error with open(config_path, "r", encoding="utf-8") as f: raw_config = json.load(f) - if not _is_qwen35_outer_config(raw_config): + if not _is_qwen35_family_outer_config(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"`config.json` for a non-Qwen3.5-family model. model={model_path} " f"error={type(error).__name__}: {error}" ) from error log_once( @@ -91,8 +91,8 @@ class QuantizationConfig: model_name: str = "qwen3_5" @classmethod - def disabled(cls) -> "QuantizationConfig": - return cls() + 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: @@ -120,7 +120,7 @@ def from_hf_config( f"{model_name} requires FP8 quantization_config; " "BF16/FP16 fallback is not supported." ) - return cls.disabled() + return cls.disabled(model_name=model_name) quant_method = str( _config_get(value, "quant_method", _config_get(value, "method", "")) @@ -216,6 +216,87 @@ def _validate_qwen35_checkpoint_precision( ) +_QWEN35_MOE_FIXED_FIELDS = { + "vocab_size": 248320, + "hidden_size": 2048, + "num_hidden_layers": 40, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "num_experts": 256, + "num_experts_per_tok": 8, + "moe_intermediate_size": 512, + "shared_expert_intermediate_size": 512, + "max_position_embeddings": 262144, +} + + +def _validate_qwen35_moe_checkpoint_config( + outer_hf_config: Any, + hf_config: Any, + quantization_config: QuantizationConfig, +) -> None: + architectures = tuple( + _config_get(outer_hf_config, "architectures", ()) or () + ) + if architectures != ("Qwen3_5MoeForConditionalGeneration",): + raise ValueError( + "Qwen3.6 MoE requires " + "architectures=['Qwen3_5MoeForConditionalGeneration'], " + f"got {list(architectures)}." + ) + if quantization_config.enabled: + raise NotImplementedError( + "Qwen3.6 MoE v1 supports BF16 checkpoints only; FP8 is out of scope." + ) + 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 NotImplementedError( + "Qwen3.6 MoE v1 requires BF16 language-model weights, " + f"got dtype={configured_dtype!r}." + ) + for field_name, expected in _QWEN35_MOE_FIXED_FIELDS.items(): + actual = _config_get(hf_config, field_name, None) + if actual != expected: + raise ValueError( + f"Qwen3.6 MoE requires {field_name}={expected!r}, " + f"got {actual!r}." + ) + expected_values = { + "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, + } + for field_name, expected in expected_values.items(): + actual = _config_get(hf_config, field_name, None) + if actual != expected: + raise ValueError( + f"Qwen3.6 MoE requires {field_name}={expected!r}, " + f"got {actual!r}." + ) + layer_types = tuple(_config_get(hf_config, "layer_types", ()) or ()) + expected_layer_types = tuple( + "full_attention" if (layer_idx + 1) % 4 == 0 else "linear_attention" + for layer_idx in range(40) + ) + if layer_types != expected_layer_types: + raise ValueError( + "Qwen3.6 MoE requires the checkpoint's 3:1 Gated DeltaNet/full-" + "attention layer layout." + ) + + _MINIMAX_M2_FIXED_FIELDS = { "vocab_size": 200064, "hidden_size": 3072, @@ -543,6 +624,17 @@ def _is_qwen35_outer_config(config: Any) -> bool: return str(_config_get(config, "model_type", "") or "").strip().lower() in {"qwen3_5", "qwen3_6"} +def _is_qwen35_moe_outer_config(config: Any) -> bool: + return str(_config_get(config, "model_type", "") or "").strip().lower() in { + "qwen3_5_moe", + "qwen3_6_moe", + } + + +def _is_qwen35_family_outer_config(config: Any) -> bool: + return _is_qwen35_outer_config(config) or _is_qwen35_moe_outer_config(config) + + def _extract_text_config(config: Any) -> Any: text_config = _config_get(config, "text_config", None) if text_config is None: @@ -678,6 +770,67 @@ def _validate_qwen3_moe_runtime(config, *, model_type: str) -> None: _validate_runtime_compatibility(config, model_type=model_type) +def _validate_qwen35_moe_runtime(config, *, model_type: str) -> None: + outer_tp_size = int(config.tensor_parallel_size) + ep_size = int(config.expert_parallel_size) + if int(config.data_parallel_size) != 1: + raise ValueError( + "Qwen3.6 MoE requires DP=1, got " + f"TP={outer_tp_size}, EP={ep_size}, DP={config.data_parallel_size}." + ) + if outer_tp_size % ep_size: + raise ValueError( + "Qwen3.6 MoE outer tensor_parallel_size must be divisible by " + f"expert_parallel_size, got outer TP={outer_tp_size}, EP={ep_size}." + ) + hf_config = config.hf_config + num_experts = int(hf_config.num_experts) + if ep_size > num_experts or num_experts % ep_size: + raise ValueError( + "Qwen3.6 MoE num_experts must be divisible by expert_parallel_size " + f"and EP must not exceed experts, got experts={num_experts}, EP={ep_size}." + ) + attention_tp_fields = { + "num_attention_heads": int(hf_config.num_attention_heads), + "num_key_value_heads": int(hf_config.num_key_value_heads), + "linear_num_key_heads": int(hf_config.linear_num_key_heads), + "linear_num_value_heads": int(hf_config.linear_num_value_heads), + "vocab_size": int(hf_config.vocab_size), + "shared_expert_intermediate_size": int( + hf_config.shared_expert_intermediate_size + ), + } + invalid_attention_fields = { + name: value + for name, value in attention_tp_fields.items() + if value % outer_tp_size + } + if invalid_attention_fields: + raise ValueError( + "Qwen3.6 MoE attention/GDN/vocabulary/shared-expert dimensions " + "must be divisible by outer tensor_parallel_size, " + f"got TP={outer_tp_size}, invalid={invalid_attention_fields}." + ) + moe_tp_size = outer_tp_size // ep_size + if int(hf_config.moe_intermediate_size) % moe_tp_size: + raise ValueError( + "Qwen3.6 MoE moe_intermediate_size must be divisible by MoE TP " + f"size, got {hf_config.moe_intermediate_size} and {moe_tp_size}." + ) + top_k = int(hf_config.num_experts_per_tok) + if not 1 <= top_k <= num_experts: + raise ValueError( + "Qwen3.6 MoE num_experts_per_tok must be in [1, num_experts], " + f"got top_k={top_k}, num_experts={num_experts}." + ) + if getattr(hf_config, "torch_dtype", None) != torch.bfloat16: + raise NotImplementedError( + "Qwen3.6 MoE v1 requires BF16 weights, got " + f"torch_dtype={getattr(hf_config, 'torch_dtype', None)}." + ) + _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) @@ -781,9 +934,13 @@ def load_and_validate_model(config) -> bool: 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) + is_qwen35_moe = _is_qwen35_moe_outer_config(config.outer_hf_config) + is_qwen35_family = is_qwen35 or is_qwen35_moe config.hf_config = _extract_text_config(config.outer_hf_config) if is_qwen35: setattr(config.hf_config, "model_type", "qwen3_5") + elif is_qwen35_moe: + setattr(config.hf_config, "model_type", "qwen3_5_moe") model_type = str(getattr(config.hf_config, "model_type", "") or "") is_minimax_m2 = model_type == "minimax_m2" is_qwen3 = model_type == "qwen3" @@ -802,8 +959,10 @@ def load_and_validate_model(config) -> bool: 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 is_qwen35_family: + raise NotImplementedError( + "Tiny random mode does not support qwen3_5 family models yet." + ) config.tiny_random_overrides = apply_tiny_random_overrides( config.hf_config, config.tiny_random_config, @@ -827,6 +986,8 @@ def load_and_validate_model(config) -> bool: quantized_model_name = "Qwen3" elif is_qwen3_moe: quantized_model_name = "Qwen3MoE" + elif is_qwen35_moe: + quantized_model_name = "Qwen3.6 MoE" config.quantization_config = QuantizationConfig.from_hf_config( raw_quantization_config, required_fp8=is_minimax_m2, @@ -838,6 +999,12 @@ def load_and_validate_model(config) -> bool: raw_quantization_config, config.quantization_config, ) + if is_qwen35_moe: + _validate_qwen35_moe_checkpoint_config( + config.outer_hf_config, + config.hf_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) @@ -865,9 +1032,11 @@ def load_and_validate_model(config) -> bool: if model_type == "qwen3_moe": _validate_qwen3_moe_runtime(config, model_type=model_type) + elif model_type == "qwen3_5_moe": + _validate_qwen35_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 + _finalize_model_config(config, is_qwen35=is_qwen35_family) + return is_qwen35_family diff --git a/src/sparsevllm/configs/runtime.py b/src/sparsevllm/configs/runtime.py index f5ffd894..38ebdb02 100644 --- a/src/sparsevllm/configs/runtime.py +++ b/src/sparsevllm/configs/runtime.py @@ -38,7 +38,7 @@ normalize_sparse_method_name, normalize_sparse_methods, ) -from sparsevllm.method_registry import PREFILL_POLICY_AUTO +from sparsevllm.method_registry import OUTER_TP_MOE_MODEL_TYPES, PREFILL_POLICY_AUTO from sparsevllm.utils.log import logger @@ -85,7 +85,7 @@ 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( + return model_type in OUTER_TP_MOE_MODEL_TYPES and int( self.tensor_parallel_size ) > 1 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..53394325 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 int(getattr(config.hf_config, "num_experts", 0) or 0) <= 0: return () max_batched_tokens = int(config.max_num_batched_tokens) @@ -574,16 +574,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..1b84f32d 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -56,6 +56,13 @@ Qwen35ForCausalLM = None _QWEN35_IMPORT_ERROR = exc +try: + from sparsevllm.models.qwen3_5_moe import Qwen35MoeForCausalLM + _QWEN35_MOE_IMPORT_ERROR = None +except ImportError as exc: + Qwen35MoeForCausalLM = None + _QWEN35_MOE_IMPORT_ERROR = exc + TP_SHM_NAME_PREFIX = "sparsevllm_" TP_SHM_SIZE = 2**20 @@ -81,6 +88,7 @@ "refresh_prefix_cache_hit", "reset_after_warmup", "run", + "runtime_diagnostic_status", "set_warmup_fake_prefill_attention", "warmup_moe_workspace", } @@ -111,6 +119,9 @@ def __init__( profiler.set_enabled(config.enable_profiler and rank == 0) hf_config = config.hf_config self.enforce_eager = config.enforce_eager + self.debug_runtime_enabled = ( + os.getenv("SPARSEVLLM_DEBUG_RUNTIME", "0") == "1" + ) self.world_size = config.world_size self.rank = rank self.event = event @@ -146,6 +157,11 @@ 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)), + ) # 加载对应的模型分片 (Shards) if hf_config.model_type == "qwen2": @@ -179,6 +195,13 @@ def __init__( 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 == "qwen3_5_moe": + if Qwen35MoeForCausalLM is None: + raise ImportError( + "Qwen35MoeForCausalLM is unavailable; verify the Qwen3.6 MoE " + f"runtime imports: {_QWEN35_MOE_IMPORT_ERROR}" + ) from _QWEN35_MOE_IMPORT_ERROR + self.model = Qwen35MoeForCausalLM(hf_config) elif hf_config.model_type == "llama": self.model = LlamaForCausalLM(hf_config) else: @@ -201,8 +224,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() @@ -822,12 +846,14 @@ def debug_sparse_state_summary(self) -> dict[str, object]: } def debug_last_logits_cpu(self) -> torch.Tensor | None: + if self.rank != 0: + return None logits = getattr(self, "debug_last_logits", None) if logits is None: raise RuntimeError( "No debug logits are available. Set SPARSEVLLM_DEBUG_RUNTIME=1 before engine startup." ) - return logits.detach().cpu() if self.rank == 0 else None + return logits.detach().cpu() def debug_hidden_states_cpu(self) -> dict[int, torch.Tensor] | None: model = getattr(getattr(self, "model", None), "model", None) @@ -911,13 +937,15 @@ def _debug_any_mismatch_from_world_rank_zero(self, tensor: torch.Tensor) -> bool def debug_replica_consistency(self) -> dict[str, object] | None: logits = getattr(self, "debug_last_logits", None) - if logits is None: + if self.world_size == 1 and logits is None: return None - logits_max_abs, logits_tolerance_ratio = self._debug_float_error_from_world_rank_zero( - logits, - atol=0.05, - rtol=0.05, - ) + if self.world_size == 1: + logits_max_abs, logits_tolerance_ratio = 0.0, 0.0 + else: + # Only world rank 0 materializes LM-head logits under tensor + # parallelism. Cross-rank consistency is checked on the synchronized + # MoE outputs below instead. + logits_max_abs, logits_tolerance_ratio = None, None result: dict[str, object] = { "last_logits_max_abs": logits_max_abs, "last_logits_tolerance_ratio": logits_tolerance_ratio, @@ -975,6 +1003,62 @@ def debug_sparse_state_summaries(self) -> list[dict[str, object]] | None: ) return summaries if self.rank == 0 else None + def runtime_diagnostic_status(self) -> list[dict[str, object]] | None: + graph_states = getattr(self.decode_cuda_graph_runner, "_graphs", {}) + graph_count = sum( + 1 + for state in graph_states.values() + if getattr(state, "graph", None) is not None + ) + local_status = { + "world_rank": int(self.parallel_context.world_rank), + "attention_tp_rank": int(self.parallel_context.attention_tp_rank), + "attention_tp_size": int(self.parallel_context.attention_tp_size), + "moe_tp_rank": int(self.parallel_context.moe_tp_rank), + "moe_tp_size": int(self.parallel_context.moe_tp_size), + "ep_rank": int(self.parallel_context.ep_rank), + "ep_size": int(self.parallel_context.ep_size), + "decode_cuda_graph_configured": bool(self.config.decode_cuda_graph), + "decode_cuda_graph_state_count": int(len(graph_states)), + "decode_cuda_graph_graph_count": int(graph_count), + "decode_cuda_graph_active": bool( + self.config.decode_cuda_graph and graph_count > 0 + ), + } + model_status = getattr(self.model, "runtime_diagnostic_status", None) + if callable(model_status): + try: + extra_status = model_status() + if not isinstance(extra_status, dict): + raise TypeError( + "Model runtime_diagnostic_status() must return a dict, " + f"got {extra_status!r}." + ) + local_status.update(extra_status) + except Exception as exc: + if self.world_size == 1: + raise + local_status["runtime_diagnostic_error"] = repr(exc) + if self.world_size == 1: + return [local_status] + statuses = [None] * self.world_size + dist.all_gather_object( + statuses, + local_status, + group=self.parallel_context.world.process_group, + ) + errors = [ + status.get("runtime_diagnostic_error") + for status in statuses + if status.get("runtime_diagnostic_error") is not None + ] + if errors: + raise RuntimeError( + "Model runtime diagnostics failed on at least one worker: " + f"{errors}." + ) + return statuses if self.rank == 0 else None + def _long_text_threshold(self, is_prefill: bool) -> int: del is_prefill if self.config.vllm_sparse_method in ("streamingllm", "attention-sink", "attention_sink"): @@ -1169,14 +1253,20 @@ def set_decode_cuda_graph_max_context_len_override(self, max_context_len: int | def set_omnikv_decode_graph_max_context_len_override(self, max_context_len: int | None): self.set_decode_cuda_graph_max_context_len_override(max_context_len) + def _capture_debug_logits(self, logits: torch.Tensor | None) -> None: + if ( + getattr(self, "debug_runtime_enabled", False) + and isinstance(logits, torch.Tensor) + ): + self.debug_last_logits = logits.detach().clone() + @torch.inference_mode() def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool): """物理执行逻辑：统一使用 Eager 模式""" _stage = 'prefill' if is_prefill else 'decode' with profiler.record(f"model_run_model_{_stage}"): logits = self.model.compute_logits(self.model(input_ids, positions)) - if os.getenv("SPARSEVLLM_DEBUG_RUNTIME", "0") == "1": - self.debug_last_logits = logits.detach().clone() + self._capture_debug_logits(logits) return logits def run_logits_for_compare(self, seqs: list[Sequence], is_prefill: bool) -> torch.Tensor | None: @@ -1257,6 +1347,7 @@ def run( else: logits = self.decode_cuda_graph_runner.run_eager_static(seqs) graph_token_ids = None + self._capture_debug_logits(logits) if self.rank != 0: self._post_sparse_forward(seqs, is_prefill) return None, None 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/method_registry.py b/src/sparsevllm/method_registry.py index 09aff5a2..2383b541 100644 --- a/src/sparsevllm/method_registry.py +++ b/src/sparsevllm/method_registry.py @@ -141,6 +141,14 @@ class ModelRuntimeCompatibility: QWEN3_MOE_TP_COMPATIBILITY = QWEN3_MOE_TP_EP_COMPATIBILITY +QWEN35_MOE_COMPATIBILITY = ModelRuntimeCompatibility( + parallel_mode="outer_tp_moe_tp_ep", + sparse_methods=frozenset({""}), + prefix_cache_methods=frozenset(), + requires_eager=False, + decode_cuda_graph_methods=frozenset({""}), +) + MINIMAX_M2_EP_COMPATIBILITY = ModelRuntimeCompatibility( parallel_mode="ep_replicated_kv", sparse_methods=frozenset( @@ -181,9 +189,20 @@ class ModelRuntimeCompatibility: MODEL_RUNTIME_COMPATIBILITY = { "qwen3_moe": QWEN3_MOE_EP_COMPATIBILITY, + "qwen3_5_moe": QWEN35_MOE_COMPATIBILITY, "minimax_m2": MINIMAX_M2_EP_COMPATIBILITY, } +OUTER_TP_MOE_MODEL_TYPES = frozenset( + {"qwen3_moe", "qwen3_5_moe", "minimax_m2"} +) + +OUTER_TP_RUNTIME_COMPATIBILITY = { + "qwen3_moe": QWEN3_MOE_TP_EP_COMPATIBILITY, + "qwen3_5_moe": QWEN35_MOE_COMPATIBILITY, + "minimax_m2": MINIMAX_M2_TP_EP_COMPATIBILITY, +} + # 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 = { @@ -266,12 +285,8 @@ def validate_model_runtime_compatibility( 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 model_type in OUTER_TP_MOE_MODEL_TYPES and tp_size > 1: + compatibility = OUTER_TP_RUNTIME_COMPATIBILITY[model_type] if dp_size != 1: raise ValueError( f"{model_type} outer_tp_moe_tp_ep requires DP=1, got " diff --git a/src/sparsevllm/models/qwen3_5.py b/src/sparsevllm/models/qwen3_5.py index 97e1c218..e7b82174 100644 --- a/src/sparsevllm/models/qwen3_5.py +++ b/src/sparsevllm/models/qwen3_5.py @@ -375,6 +375,16 @@ 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) @@ -625,7 +635,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 +673,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,17 +823,20 @@ 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) -> 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, @@ -845,7 +860,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 +872,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 +896,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 diff --git a/src/sparsevllm/models/qwen3_5_moe.py b/src/sparsevllm/models/qwen3_5_moe.py new file mode 100644 index 00000000..ce0cc121 --- /dev/null +++ b/src/sparsevllm/models/qwen3_5_moe.py @@ -0,0 +1,542 @@ +from __future__ import annotations + +import os +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.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$" +) + + +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, self.hidden_size) + ) + + def forward( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + router_logits = F.linear(hidden_states, self.weight) + topk_weights, topk_ids = self.provider.run( + self.op_spec, router_logits + ) + return router_logits, topk_weights, topk_ids + + +class Qwen35MoePackedExperts(Qwen3MoePackedExperts): + """Qwen3.6 packed-3D checkpoint adapter for the shared MoE provider.""" + + checkpoint_projection_map = { + "gate_up_proj": "gate_up", + "down_proj": "down", + } + + def __init__(self, config) -> None: + super().__init__(config) + if self.fp8_enabled: + raise NotImplementedError( + "Qwen3.6 MoE v1 accepts BF16 expert weights only." + ) + self._loaded_packed_projections: set[str] = set() + + def rank_local_weight_slice( + self, + source_shape: tuple[int, ...], + *, + loaded_shard_id: str, + is_scale: bool = False, + ) -> tuple[slice, ...] | None: + 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: + 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: + 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.debug_enabled = os.getenv("SPARSEVLLM_DEBUG_MOE", "0") == "1" + 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), + ) + self.shared_expert_gate = nn.Linear( + int(config.hidden_size), 1, bias=False + ) + + def _forward_chunk( + self, + hidden_states: torch.Tensor, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + shared_output = self.shared_expert(hidden_states) + router_logits, topk_weights, topk_ids = self.gate(hidden_states) + local_output = self.experts(hidden_states, topk_ids, topk_weights) + routed_output = self.parallel_context.world_all_reduce(local_output) + shared_gate = torch.sigmoid(self.shared_expert_gate(hidden_states)) + gated_shared_output = shared_gate * shared_output + return ( + routed_output + gated_shared_output, + router_logits, + topk_weights, + topk_ids, + local_output, + gated_shared_output, + ) + + 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)}." + ) + debug_enabled = self.debug_enabled + if debug_enabled: + self.debug_last_input = hidden_states.detach().clone() + + chunks = hidden_states.split(self.mlp_chunk_size, dim=0) + outputs = [] + router_logits_chunks = [] + topk_weights_chunks = [] + topk_ids_chunks = [] + local_output_chunks = [] + shared_output_chunks = [] + for chunk in chunks: + ( + output, + router_logits, + topk_weights, + topk_ids, + local_output, + shared_output, + ) = self._forward_chunk(chunk) + outputs.append(output) + if debug_enabled: + router_logits_chunks.append(router_logits) + topk_weights_chunks.append(topk_weights) + topk_ids_chunks.append(topk_ids) + local_output_chunks.append(local_output) + shared_output_chunks.append(shared_output) + output = outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0) + + if debug_enabled: + self.debug_last_router_logits = torch.cat( + router_logits_chunks, dim=0 + ).detach().clone() + self.debug_last_topk_weights = torch.cat( + topk_weights_chunks, dim=0 + ).detach().clone() + self.debug_last_topk_ids = torch.cat( + topk_ids_chunks, dim=0 + ).detach().clone() + self.debug_last_local_output = torch.cat( + local_output_chunks, dim=0 + ).detach().clone() + self.debug_last_shared_output = torch.cat( + shared_output_chunks, dim=0 + ).detach().clone() + local_mask = ( + self.debug_last_topk_ids >= self.experts.local_expert_start + ) & (self.debug_last_topk_ids < self.experts.local_expert_end) + local_hit_count = local_mask.sum() + self.debug_last_local_hit_count = ( + local_hit_count + if device_runtime.is_stream_capturing() + else int(local_hit_count.item()) + ) + self.debug_last_output = output.detach().clone() + return output + + +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", + } + + 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() + + @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, + ), + ), + ) + + def runtime_diagnostic_status(self) -> dict[str, object]: + experts = self.model.layers[0].mlp.experts + router = self.model.layers[0].mlp.gate + return { + "moe_expert_provider": experts.provider.name, + "moe_router_provider": router.provider.name, + "local_expert_start": int(experts.local_expert_start), + "local_expert_end": int(experts.local_expert_end), + } + + @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: + match = _PACKED_EXPERT_SOURCE_RE.match(source_weight_name) + if match is not None: + layer_idx, projection = match.groups() + 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 = _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 = _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: + del loaded_weight_shape, loaded_weight_dtype + if not source_weight_name.startswith(self.ignored_weight_prefixes): + raise ValueError( + f"Qwen3.6 MoE loader unexpectedly skipped {source_weight_name!r}." + ) + if loaded_scale_shape is not None or loaded_scale_dtype is not None: + raise ValueError( + "Qwen3.6 MoE visual/MTP intentional skips must not consume " + f"quantization scales: {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() + + 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}." + ) + first_experts = self.model.layers[0].mlp.experts + logger.info( + "Loaded Qwen3.6 MoE rank {} expert_provider={} router_provider={} " + "attention TP {}/{} " + "MoE TP {}/{} EP {}/{} local experts [{}, {}) across {} layers; " + "intentionally skipped {} visual/MTP tensors.", + self.parallel_context.world_rank, + 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/operators/moe.py b/src/sparsevllm/operators/moe.py index e0d212fd..d9530acd 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -1,9 +1,11 @@ from __future__ import annotations +import os from dataclasses import dataclass from importlib.util import find_spec import torch +import torch.nn.functional as F import sparsevllm.platforms as platforms from sparsevllm.operators.registry import ( @@ -176,6 +178,69 @@ def run( MOE_REGISTRY: OpRegistry[MoeOpSpec, MoeProvider] = OpRegistry("routed MoE") +@MOE_REGISTRY.register +class TorchMoeProvider(MoeProvider): + """Clear eager-only CUDA reference for routed-expert semantics.""" + + name = "torch" + priority = 0 + gate_up_order = "gate_up" + + @classmethod + def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: + if caps.platform != PlatformEnum.CUDA: + return SupportResult.no(f"requires CUDA, got {caps.platform.name}") + if spec.cuda_graph: + return SupportResult.no("reference provider is eager-only") + if spec.weight_dtype != spec.activation_dtype: + return SupportResult.no( + "reference provider requires unquantized weights matching activations" + ) + if spec.block_shape is not None: + return SupportResult.no("reference provider does not support quantized weights") + 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, + ): + del ep_rank + if w13_scale_inv is not None or w2_scale_inv is not None: + raise RuntimeError("Torch MoE reference does not accept expert scales.") + output = torch.zeros_like(hidden_states) + for local_expert_id in range(spec.num_local_experts): + global_expert_id = int(local_expert_start) + local_expert_id + token_indices, topk_slots = torch.where( + topk_ids == global_expert_id + ) + if token_indices.numel() == 0: + continue + projected = F.linear( + hidden_states[token_indices], + w13_weight[local_expert_id], + ) + gate, up = projected.chunk(2, dim=-1) + expert_output = F.linear( + F.silu(gate) * up, + w2_weight[local_expert_id], + ) + routed = expert_output * topk_weights[ + token_indices, topk_slots + ].unsqueeze(-1) + output.index_add_(0, token_indices, routed) + return output + + @MOE_REGISTRY.register class TritonMinimaxM2FusedMoeProvider(MoeProvider): name = "triton_minimax_m2_fused" @@ -513,4 +578,22 @@ def resolve_moe_provider( 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_REGISTRY).resolve(spec, caps).provider + requested = os.getenv("SPARSEVLLM_MOE_PROVIDER", "auto").strip().lower() + if requested == "auto": + return OpResolver(MOE_REGISTRY).resolve(spec, caps).provider + + providers = {provider.name: provider for provider in MOE_REGISTRY.providers} + if requested not in providers: + choices = ", ".join(["auto", *sorted(providers)]) + raise ValueError( + "SPARSEVLLM_MOE_PROVIDER must be one of " + f"{choices}, got {requested!r}." + ) + provider_cls = providers[requested] + support = provider_cls.supports(spec, caps) + if not support.supported: + raise RuntimeError( + f"Requested MoE provider {requested!r} does not support " + f"spec={spec!r} on device={caps.device_name!r}: {support.reason}." + ) + return provider_cls() diff --git a/src/sparsevllm/operators/moe_router.py b/src/sparsevllm/operators/moe_router.py new file mode 100644 index 00000000..57334719 --- /dev/null +++ b/src/sparsevllm/operators/moe_router.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import os +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, + ) + + +@MOE_ROUTER_REGISTRY.register +class TorchMoeRouterProvider(MoeRouterProvider): + name = "torch" + priority = 0 + + @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 spec.cuda_graph: + return SupportResult.no("reference provider is eager-only") + return SupportResult.yes() + + def run( + self, + spec: MoeRouterOpSpec, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + probabilities = torch.softmax(router_logits, dim=-1, dtype=torch.float32) + weights, ids = torch.topk(probabilities, spec.top_k, dim=-1) + if spec.norm_topk_prob: + weights = weights / weights.sum(dim=-1, keepdim=True) + return weights.to(dtype=router_logits.dtype), ids.to(dtype=torch.int32) + + +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)) + requested = os.getenv("SPARSEVLLM_MOE_ROUTER_PROVIDER", "auto").strip().lower() + if requested == "auto": + return OpResolver(MOE_ROUTER_REGISTRY).resolve(spec, caps).provider + + providers = { + provider.name: provider for provider in MOE_ROUTER_REGISTRY.providers + } + if requested not in providers: + choices = ", ".join(["auto", *sorted(providers)]) + raise ValueError( + "SPARSEVLLM_MOE_ROUTER_PROVIDER must be one of " + f"{choices}, got {requested!r}." + ) + provider_cls = providers[requested] + support = provider_cls.supports(spec, caps) + if not support.supported: + raise RuntimeError( + f"Requested MoE router provider {requested!r} does not support " + f"spec={spec!r} on device={caps.device_name!r}: {support.reason}." + ) + return provider_cls() diff --git a/src/sparsevllm/triton_kernel/moe_topk.py b/src/sparsevllm/triton_kernel/moe_topk.py index 5c115d67..0bdc2a4a 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, @@ -200,9 +202,11 @@ def topk_softmax( raise ValueError("Triton topk_softmax requires contiguous router_logits.") 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 in {128, 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/tests/test_input_processor.py b/tests/test_input_processor.py new file mode 100644 index 00000000..d78f4dc3 --- /dev/null +++ b/tests/test_input_processor.py @@ -0,0 +1,54 @@ +from unittest.mock import Mock + +import pytest + +from sparsevllm.engine.input_processor import tokenize_text_prompt + + +def test_tokenize_text_prompt_preserves_token_ids(): + tokenizer = Mock() + + assert tokenize_text_prompt(tokenizer, [1, 2, 3]) == [1, 2, 3] + tokenizer.encode.assert_not_called() + + +@pytest.mark.parametrize( + "prompt", + [ + {"image": "example.png", "text": "describe"}, + {"video": "example.mp4", "text": "describe"}, + {"mtp": True, "text": "continue"}, + [1, "2"], + ], +) +def test_tokenize_text_prompt_rejects_structured_inputs(prompt): + with pytest.raises(TypeError, match="text prompts only"): + tokenize_text_prompt(Mock(), prompt) + + +def test_tokenize_text_prompt_rejects_empty_token_ids(): + with pytest.raises(ValueError, match="must not be empty"): + tokenize_text_prompt(Mock(), []) + + +@pytest.mark.parametrize( + ("bos_token", "prompt", "add_special_tokens"), + [ + ("", "hello", True), + ("", "hello", False), + (None, "hello", False), + ], +) +def test_tokenize_text_prompt_handles_bos_once( + bos_token, + prompt, + add_special_tokens, +): + tokenizer = Mock(bos_token=bos_token) + tokenizer.encode.return_value = [1] + + assert tokenize_text_prompt(tokenizer, prompt) == [1] + tokenizer.encode.assert_called_once_with( + prompt, + add_special_tokens=add_special_tokens, + ) diff --git a/tests/test_longbench_deltakv_contracts.py b/tests/test_longbench_deltakv_contracts.py index 1941209d..eaba245e 100644 --- a/tests/test_longbench_deltakv_contracts.py +++ b/tests/test_longbench_deltakv_contracts.py @@ -60,6 +60,48 @@ def test_longbench_data_validation_requires_explicit_root(self): finally: longbench_pred.DATA_PREFIX_PATH = old_root + def test_longbench_writes_prompt_and_canonical_per_sample_artifact(self): + with tempfile.TemporaryDirectory() as tmp: + task_path = str(Path(tmp) / "qasper.jsonl") + record = { + "dataset": "qasper", + "sample_idx": 0, + "source_idx": 12, + "status": "success", + "prompt_tokens": 16001, + "rendered_prompt": "fixed rendered prompt", + "rendered_prompt_sha256": longbench_pred._sha256_text( + "fixed rendered prompt" + ), + "pred": "answer", + "raw_pred": "answer", + "answers": ["answer"], + "all_classes": [], + "length": 16001, + } + + longbench_pred._write_sample_record( + out_root=tmp, + task_out_path=task_path, + record=record, + ) + + canonical = (Path(tmp) / "per_sample_results.jsonl").read_text( + encoding="utf-8" + ) + historical = (Path(tmp) / "sample_results.jsonl").read_text( + encoding="utf-8" + ) + raw = json.loads( + (Path(tmp) / "raw_outputs.jsonl").read_text(encoding="utf-8") + ) + self.assertEqual(canonical, historical) + self.assertEqual(raw["rendered_prompt"], "fixed rendered prompt") + self.assertEqual( + raw["rendered_prompt_sha256"], + longbench_pred._sha256_text("fixed rendered prompt"), + ) + def test_sparsevllm_data_workers_receive_distinct_master_ports(self): launched = [] @@ -95,18 +137,27 @@ def fake_popen(command, *, env, cwd): ) def test_longbench_records_actual_decode_cuda_graph_state(self): - graph_runner = SimpleNamespace( - _graphs={ - "captured": SimpleNamespace(graph=object()), - "uncaptured": SimpleNamespace(graph=None), + worker_statuses = [ + { + "world_rank": 0, + "decode_cuda_graph_configured": True, + "decode_cuda_graph_state_count": 2, + "decode_cuda_graph_graph_count": 1, + "decode_cuda_graph_active": True, }, - last_state_key="captured", - ) + { + "world_rank": 1, + "decode_cuda_graph_configured": True, + "decode_cuda_graph_state_count": 2, + "decode_cuda_graph_graph_count": 1, + "decode_cuda_graph_active": True, + }, + ] generate_fn = SimpleNamespace( _sparsevllm_llm=SimpleNamespace( - config=SimpleNamespace(decode_cuda_graph=True), + config=SimpleNamespace(world_size=2), model_runner=SimpleNamespace( - decode_cuda_graph_runner=graph_runner, + call=lambda method: worker_statuses ), ) ) @@ -119,18 +170,71 @@ def test_longbench_records_actual_decode_cuda_graph_state(self): ) path = Path(tmp) / "decode_cuda_graph_status_rank2.json" - self.assertEqual(status["rank"], 2) - self.assertTrue(status["configured"]) - self.assertTrue(status["runner_initialized"]) - self.assertEqual(status["state_count"], 2) - self.assertEqual(status["graph_count"], 1) - self.assertTrue(status["active"]) - self.assertEqual(status["last_state_key"], "captured") + self.assertEqual(status["launcher_rank"], 2) + self.assertTrue(status["configured_on_all_workers"]) + self.assertTrue(status["active_on_all_workers"]) + self.assertEqual(status["workers"], worker_statuses) self.assertEqual( json.loads(path.read_text(encoding="utf-8")), status, ) + def test_longbench_fails_if_any_requested_graph_worker_is_inactive(self): + generate_fn = SimpleNamespace( + _sparsevllm_llm=SimpleNamespace( + config=SimpleNamespace(world_size=2), + model_runner=SimpleNamespace( + call=lambda method: [ + { + "world_rank": 0, + "decode_cuda_graph_configured": True, + "decode_cuda_graph_active": True, + }, + { + "world_rank": 1, + "decode_cuda_graph_configured": True, + "decode_cuda_graph_active": False, + }, + ] + ), + ) + ) + + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaisesRegex(RuntimeError, "not active on every"): + longbench_pred._write_decode_cuda_graph_status( + generate_fn=generate_fn, + out_root=tmp, + rank=0, + ) + + def test_longbench_fails_if_requested_graph_is_unconfigured(self): + generate_fn = SimpleNamespace( + _sparsevllm_llm=SimpleNamespace( + config=SimpleNamespace( + world_size=1, + decode_cuda_graph=True, + ), + model_runner=SimpleNamespace( + call=lambda method: [ + { + "world_rank": 0, + "decode_cuda_graph_configured": False, + "decode_cuda_graph_active": False, + } + ] + ), + ) + ) + + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaisesRegex(RuntimeError, "not configured on every"): + longbench_pred._write_decode_cuda_graph_status( + generate_fn=generate_fn, + out_root=tmp, + rank=0, + ) + def test_longbench_fails_if_sparsevllm_graph_state_is_unavailable(self): with tempfile.TemporaryDirectory() as tmp: with self.assertRaisesRegex(RuntimeError, "_sparsevllm_llm"): diff --git a/tests/test_microbench_artifacts.py b/tests/test_microbench_artifacts.py index 3bdc0c99..3f4f2b7d 100644 --- a/tests/test_microbench_artifacts.py +++ b/tests/test_microbench_artifacts.py @@ -9,10 +9,32 @@ _benchmark_sparse_method, _record_child_exit_failure, _resolved_engine_config, + _worker_runtime_status, _write_output_dir, ) +def test_worker_runtime_status_preserves_per_rank_graph_and_provider_data(): + expected = [ + { + "world_rank": 0, + "decode_cuda_graph_active": True, + "moe_router_provider": "triton", + }, + { + "world_rank": 1, + "decode_cuda_graph_active": True, + "moe_router_provider": "triton", + }, + ] + llm = SimpleNamespace( + config=SimpleNamespace(world_size=2), + model_runner=SimpleNamespace(call=lambda method: expected), + ) + + assert _worker_runtime_status(llm) == expected + + @pytest.mark.parametrize( ("method", "expected"), [ diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 6e05b2e4..81ca023b 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -1,3 +1,4 @@ +import os import sys from types import SimpleNamespace from unittest.mock import Mock, patch @@ -12,7 +13,16 @@ TritonFp8LinearProvider, resolve_fp8_linear_provider, ) -from sparsevllm.operators.moe import MOE_REGISTRY, MoeOpSpec, resolve_moe_provider +from sparsevllm.operators.moe import ( + MOE_REGISTRY, + MoeOpSpec, + TorchMoeProvider, + resolve_moe_provider, +) +from sparsevllm.operators.moe_router import ( + MoeRouterOpSpec, + resolve_moe_router_provider, +) from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum @@ -64,6 +74,7 @@ def _moe_spec( tp_size=1, routing_method="softmax", scale_dtype=None, + cuda_graph=True, ) -> MoeOpSpec: return MoeOpSpec( num_experts=num_experts, @@ -75,7 +86,7 @@ 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, @@ -346,6 +357,110 @@ def test_unquantized_moe_uses_triton_on_supported_cuda(dtype, capability): assert resolved.provider.name == "triton" +def test_torch_moe_reference_matches_explicit_expert_routing(): + spec = _moe_spec( + activation_dtype=torch.float32, + weight_dtype=torch.float32, + block_shape=None, + hidden_size=3, + intermediate_size=2, + num_local_experts=2, + num_experts=4, + top_k=2, + ep_size=2, + cuda_graph=False, + ) + hidden_states = torch.tensor( + [[1.0, -2.0, 0.5], [0.25, 1.0, -0.75], [-1.0, 0.5, 2.0]] + ) + topk_ids = torch.tensor([[2, 0], [3, 2], [1, 0]]) + topk_weights = torch.tensor([[0.7, 0.3], [0.4, 0.6], [0.8, 0.2]]) + w13_weight = torch.arange(24, dtype=torch.float32).reshape(2, 4, 3) / 17 + w2_weight = torch.arange(12, dtype=torch.float32).reshape(2, 3, 2) / 11 + + actual = TorchMoeProvider().run( + spec, + hidden_states, + topk_ids, + topk_weights, + w13_weight, + w2_weight, + None, + None, + local_expert_start=2, + ep_rank=1, + ) + + expected = torch.zeros_like(hidden_states) + for token_index, routes in enumerate(zip(topk_ids, topk_weights)): + for expert_id, route_weight in zip(*routes): + local_expert_id = int(expert_id) - 2 + if not 0 <= local_expert_id < 2: + continue + gate_up = hidden_states[token_index] @ w13_weight[local_expert_id].T + gate, up = gate_up.chunk(2) + expert_output = ( + torch.nn.functional.silu(gate) * up + ) @ w2_weight[local_expert_id].T + expected[token_index] += expert_output * route_weight + + torch.testing.assert_close(actual, expected) + assert torch.equal(actual[2], torch.zeros(3)) + + +def test_explicit_torch_moe_provider_selection_is_eager_only(): + caps = _cuda_caps((9, 0), native_fp8=False) + platform = SimpleNamespace(get_device_caps=lambda _: caps) + eager_spec = _moe_spec( + activation_dtype=torch.bfloat16, + weight_dtype=torch.bfloat16, + block_shape=None, + cuda_graph=False, + ) + graph_spec = _moe_spec( + activation_dtype=torch.bfloat16, + weight_dtype=torch.bfloat16, + block_shape=None, + cuda_graph=True, + ) + + with ( + patch("sparsevllm.operators.moe.platforms.current_platform", platform), + patch.dict(os.environ, {"SPARSEVLLM_MOE_PROVIDER": "torch"}), + ): + assert resolve_moe_provider(eager_spec, device_index=0).name == "torch" + with pytest.raises(RuntimeError, match="eager-only"): + resolve_moe_provider(graph_spec, device_index=0) + + +def test_explicit_torch_router_provider_selection_is_eager_only(): + caps = _cuda_caps((9, 0), native_fp8=False) + platform = SimpleNamespace(get_device_caps=lambda _: caps) + + def spec(cuda_graph): + return MoeRouterOpSpec( + num_experts=256, + top_k=8, + activation_dtype=torch.bfloat16, + norm_topk_prob=True, + cuda_graph=cuda_graph, + ) + + with ( + patch( + "sparsevllm.operators.moe_router.platforms.current_platform", + platform, + ), + patch.dict( + os.environ, + {"SPARSEVLLM_MOE_ROUTER_PROVIDER": "torch"}, + ), + ): + assert resolve_moe_router_provider(spec(False), device_index=0).name == "torch" + with pytest.raises(RuntimeError, match="eager-only"): + resolve_moe_router_provider(spec(True), device_index=0) + + def test_hopper_fused_moe_uses_profiled_tp_ep_shape(): spec = _moe_spec( activation_dtype=torch.bfloat16, diff --git a/tests/test_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index f1f44396..43b03bdf 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -1685,7 +1685,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.hf_config = SimpleNamespace(model_type="qwen3_moe", num_experts=128) self.assertEqual( _moe_workspace_warmup_token_counts(config), @@ -1704,7 +1704,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"), + hf_config=SimpleNamespace(model_type="qwen3_moe", num_experts=128), ) calls = [] engine.model_runner = SimpleNamespace( @@ -1727,7 +1727,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"), + hf_config=SimpleNamespace(model_type="qwen3_moe", num_experts=128), ) 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..5bc3a2c5 100644 --- a/tests/test_qwen35_mixed_runtime.py +++ b/tests/test_qwen35_mixed_runtime.py @@ -54,6 +54,37 @@ def _single_process_parallel_context() -> ParallelContext: return ParallelContext(world=group, tensor=group, expert=group, data=group) +def test_model_runner_debug_logits_capture_is_explicit(): + runner = SimpleNamespace(debug_runtime_enabled=True) + + ModelRunner._capture_debug_logits(runner, None) + assert not hasattr(runner, "debug_last_logits") + + logits = torch.tensor([[1.0, 2.0]]) + ModelRunner._capture_debug_logits(runner, logits) + torch.testing.assert_close(runner.debug_last_logits, logits) + assert runner.debug_last_logits.data_ptr() != logits.data_ptr() + + +def test_model_runner_debug_logits_non_output_rank_returns_none(): + runner = SimpleNamespace(rank=1) + + assert ModelRunner.debug_last_logits_cpu(runner) is None + + +def test_model_runner_tp_replica_consistency_does_not_require_rank_local_logits(): + runner = SimpleNamespace( + world_size=2, + model=SimpleNamespace(model=SimpleNamespace(layers=[])), + ) + + assert ModelRunner.debug_replica_consistency(runner) == { + "last_logits_max_abs": None, + "last_logits_tolerance_ratio": None, + "moe_layers": {}, + } + + def _qwen35_outer_config(*, num_layers: int = 64, full_layers: tuple[int, ...] | None = None): if full_layers is None: full_layers = tuple(range(0, num_layers, 4)) diff --git a/tests/test_qwen35_moe.py b/tests/test_qwen35_moe.py new file mode 100644 index 00000000..d370bd4a --- /dev/null +++ b/tests/test_qwen35_moe.py @@ -0,0 +1,403 @@ +from contextlib import ExitStack +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from sparsevllm.config import Config, QuantizationConfig +from sparsevllm.distributed import ParallelContext, ParallelGroup +from sparsevllm.engine.recurrent_state_manager import RecurrentStateManager +from sparsevllm.models.qwen3_5_moe import ( + Qwen35MoeForCausalLM, + Qwen35MoePackedExperts, + Qwen35MoeSparseMoeBlock, +) +from sparsevllm.models.qwen3_5 import Qwen35LinearAttention +from sparsevllm.operators.moe import TritonMoeProvider +from sparsevllm.operators.moe_router import ( + MoeRouterOpSpec, + TorchMoeRouterProvider, +) + + +def _outer_config(): + layer_types = [ + "full_attention" if (layer_idx + 1) % 4 == 0 else "linear_attention" + for layer_idx in range(40) + ] + text_config = SimpleNamespace( + model_type="qwen3_5_moe_text", + vocab_size=248320, + hidden_size=2048, + num_hidden_layers=40, + layer_types=layer_types, + num_attention_heads=16, + num_key_value_heads=2, + head_dim=256, + linear_num_key_heads=16, + linear_num_value_heads=32, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + moe_intermediate_size=512, + shared_expert_intermediate_size=512, + num_experts=256, + num_experts_per_tok=8, + hidden_act="silu", + attn_output_gate=True, + attention_bias=False, + partial_rotary_factor=0.25, + mamba_ssm_dtype="float32", + tie_word_embeddings=False, + rms_norm_eps=1.0e-6, + max_position_embeddings=262144, + torch_dtype=torch.bfloat16, + quantization_config=None, + ) + return SimpleNamespace( + model_type="qwen3_5_moe", + architectures=["Qwen3_5MoeForConditionalGeneration"], + text_config=text_config, + ) + + +def _make_config(tmp_path, **overrides): + with patch( + "sparsevllm.configs.runtime.AutoConfig.from_pretrained", + return_value=_outer_config(), + ): + return Config(model=str(tmp_path), **overrides) + + +def _hybrid_context(world_rank: int) -> ParallelContext: + world_ranks = (0, 1, 2, 3) + moe_tp_ranks = (0, 1) if world_rank < 2 else (2, 3) + moe_ep_ranks = (0, 2) if world_rank % 2 == 0 else (1, 3) + return ParallelContext( + world=ParallelGroup(None, world_ranks, world_rank, 4), + tensor=ParallelGroup(None, world_ranks, world_rank, 4), + expert=ParallelGroup( + None, moe_ep_ranks, moe_ep_ranks.index(world_rank), 2 + ), + data=ParallelGroup(None, (world_rank,), 0, 1), + moe_tensor=ParallelGroup( + None, moe_tp_ranks, moe_tp_ranks.index(world_rank), 2 + ), + ) + + +def _single_context() -> ParallelContext: + group = ParallelGroup(None, (0,), 0, 1) + return ParallelContext(group, group, group, group) + + +def _pure_tp_context(world_rank: int) -> ParallelContext: + world = ParallelGroup(None, (0, 1), world_rank, 2) + singleton = ParallelGroup(None, (world_rank,), 0, 1) + return ParallelContext( + world=world, + tensor=world, + expert=singleton, + data=singleton, + moe_tensor=world, + ) + + +def test_qwen36_moe_config_normalizes_text_runtime_and_topology(tmp_path): + config = _make_config( + tmp_path, + tensor_parallel_size=2, + expert_parallel_size=2, + decode_cuda_graph=True, + enforce_eager=False, + ) + + assert config.hf_config.model_type == "qwen3_5_moe" + assert config.uses_outer_tp_moe_layout is True + assert config.world_size == 2 + assert config.moe_tensor_parallel_size == 1 + assert config.runtime_layout.full_attention_layer_indices == tuple( + range(3, 40, 4) + ) + assert config.runtime_layout.num_kv_layers == 10 + + +def test_qwen36_moe_rejects_non_vanilla_sparse_method(tmp_path): + with pytest.raises(ValueError, match="validated methods: 'vanilla'"): + _make_config(tmp_path, vllm_sparse_method="quest") + + +def test_qwen36_moe_rejects_invalid_outer_tp_ep_topology(tmp_path): + with pytest.raises(ValueError, match="must be divisible"): + _make_config( + tmp_path, + tensor_parallel_size=2, + expert_parallel_size=3, + ) + + +def test_qwen36_moe_rejects_non_bf16_checkpoint(tmp_path): + outer = _outer_config() + outer.text_config.torch_dtype = torch.float16 + with patch( + "sparsevllm.configs.runtime.AutoConfig.from_pretrained", + return_value=outer, + ): + with pytest.raises(NotImplementedError, match="requires BF16"): + Config(model=str(tmp_path)) + + +def test_qwen36_moe_recurrent_state_uses_attention_tp_and_fp32_state(): + config = _outer_config().text_config + + spec = Qwen35MoeForCausalLM.recurrent_state_spec( + config, attention_tp_size=2 + ) + + assert spec.tensor_specs[0].shape == (4096, 3) + assert spec.tensor_specs[0].dtype == torch.bfloat16 + assert spec.tensor_specs[1].shape == (16, 128, 128) + assert spec.tensor_specs[1].dtype == torch.float32 + + +def test_qwen36_moe_torch_router_is_fp32_softmax_oracle(): + logits = torch.tensor( + [[-80.0, -2.0, 0.0, 1.0, 3.0, 7.0, 8.0, 9.0, 10.0]], + dtype=torch.bfloat16, + ) + spec = MoeRouterOpSpec( + num_experts=9, + top_k=8, + activation_dtype=torch.bfloat16, + norm_topk_prob=True, + cuda_graph=False, + ) + + weights, ids = TorchMoeRouterProvider().run(spec, logits) + probabilities = torch.softmax(logits, dim=-1, dtype=torch.float32) + expected_weights, expected_ids = torch.topk(probabilities, 8, dim=-1) + expected_weights /= expected_weights.sum(dim=-1, keepdim=True) + + assert weights.dtype == torch.bfloat16 + assert ids.dtype == torch.int32 + assert torch.equal(ids, expected_ids.to(torch.int32)) + torch.testing.assert_close( + weights.float(), expected_weights, atol=4e-3, rtol=4e-3 + ) + + +def test_qwen36_moe_linear_attention_uses_configured_recurrent_dtype(): + config = SimpleNamespace( + hidden_size=8, + hidden_act="silu", + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=4, + linear_value_head_dim=4, + linear_conv_kernel_dim=4, + rms_norm_eps=1.0e-6, + mlp_chunk_size=16, + torch_dtype=torch.bfloat16, + runtime_recurrent_state_dtype=torch.float32, + quantization_config=QuantizationConfig.disabled(), + ) + context = _single_context() + with ( + patch( + "sparsevllm.models.qwen3_5.get_parallel_context", + return_value=context, + ), + patch( + "sparsevllm.layers.linear.get_parallel_context", + return_value=context, + ), + ): + attention = Qwen35LinearAttention(config) + + assert attention.recurrent_state_dtype == torch.float32 + + +def test_recurrent_pool_accepts_model_declared_mixed_state_dtypes(): + runtime_config = SimpleNamespace( + runtime_layout=SimpleNamespace( + linear_attention_layer_indices=(0,), + is_linear_attention=lambda layer_idx: int(layer_idx) == 0, + ), + enable_prefix_caching=False, + max_num_seqs_in_batch=1, + max_decoding_seqs=1, + max_num_seqs_in_gpu=1, + recurrent_state_max_bytes=None, + prefix_cache_block_size=4, + ) + state_spec = Qwen35MoeForCausalLM.recurrent_state_spec( + _outer_config().text_config, + attention_tp_size=1, + ) + manager = RecurrentStateManager( + runtime_config, + _single_context(), + device=torch.device("cpu"), + state_spec=state_spec, + ) + seq = SimpleNamespace(seq_id=1) + manager.prepare_step([seq], is_prefill=False) + manager.prepare_decode_static([seq], token_batch=1, device=torch.device("cpu")) + + state_buffers, _ = manager.get_decode_layer_state( + [seq], + layer_idx=0, + token_batch=1, + dtype=torch.bfloat16, + device=torch.device("cpu"), + ) + + assert state_buffers["conv_state"].dtype == torch.bfloat16 + assert state_buffers["recurrent_state"].dtype == torch.float32 + + +def test_packed_experts_slice_ep_before_moe_tp(): + context = _hybrid_context(world_rank=1) + config = SimpleNamespace( + num_experts=4, + hidden_size=4, + moe_intermediate_size=4, + num_experts_per_tok=2, + dtype=torch.bfloat16, + torch_dtype=torch.bfloat16, + decode_cuda_graph=False, + quantization_config=QuantizationConfig.disabled(), + ) + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + with ExitStack() as stack: + stack.enter_context( + patch( + "sparsevllm.models.qwen3_moe.get_parallel_context", + return_value=context, + ) + ) + stack.enter_context( + patch( + "sparsevllm.models.qwen3_moe.resolve_moe_provider", + return_value=TritonMoeProvider(), + ) + ) + experts = Qwen35MoePackedExperts(config) + finally: + torch.set_default_dtype(previous_dtype) + + gate_up_global_shape = (4, 8, 4) + down_global_shape = (4, 4, 4) + assert experts.rank_local_weight_slice( + gate_up_global_shape, + loaded_shard_id="gate_up_proj", + ) == (slice(0, 2), slice(None), slice(None)) + assert experts.rank_local_weight_slice( + down_global_shape, + loaded_shard_id="down_proj", + ) == (slice(0, 2), slice(None), slice(None)) + + gate_up = torch.arange(2 * 8 * 4, dtype=torch.bfloat16).view(2, 8, 4) + down = torch.arange(2 * 4 * 4, dtype=torch.bfloat16).view(2, 4, 4) + experts.load_packed_expert_weight("gate_up_proj", gate_up) + experts.load_packed_expert_weight("down_proj", down) + experts.validate_loaded_weights() + + expected_gate = gate_up[:, 2:4] + expected_up = gate_up[:, 6:8] + torch.testing.assert_close(experts.w13_weight[:, :2], expected_gate) + torch.testing.assert_close(experts.w13_weight[:, 2:], expected_up) + torch.testing.assert_close(experts.w2_weight, down[:, :, 2:4]) + + +def test_packed_expert_pure_tp_shards_reconstruct_checkpoint(): + config = SimpleNamespace( + num_experts=4, + hidden_size=4, + moe_intermediate_size=4, + num_experts_per_tok=2, + dtype=torch.bfloat16, + torch_dtype=torch.bfloat16, + decode_cuda_graph=False, + quantization_config=QuantizationConfig.disabled(), + ) + gate_up = torch.arange(4 * 8 * 4, dtype=torch.bfloat16).view(4, 8, 4) + down = torch.arange(4 * 4 * 4, dtype=torch.bfloat16).view(4, 4, 4) + rank_experts = [] + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + for world_rank in (0, 1): + with ExitStack() as stack: + stack.enter_context( + patch( + "sparsevllm.models.qwen3_moe.get_parallel_context", + return_value=_pure_tp_context(world_rank), + ) + ) + stack.enter_context( + patch( + "sparsevllm.models.qwen3_moe.resolve_moe_provider", + return_value=TritonMoeProvider(), + ) + ) + experts = Qwen35MoePackedExperts(config) + experts.load_packed_expert_weight("gate_up_proj", gate_up) + experts.load_packed_expert_weight("down_proj", down) + experts.validate_loaded_weights() + rank_experts.append(experts) + finally: + torch.set_default_dtype(previous_dtype) + + gate = torch.cat( + [experts.w13_weight[:, :2] for experts in rank_experts], dim=1 + ) + up = torch.cat( + [experts.w13_weight[:, 2:] for experts in rank_experts], dim=1 + ) + reconstructed_down = torch.cat( + [experts.w2_weight for experts in rank_experts], dim=2 + ) + torch.testing.assert_close(gate, gate_up[:, :4]) + torch.testing.assert_close(up, gate_up[:, 4:]) + torch.testing.assert_close(reconstructed_down, down) + + +def test_routed_output_reduces_without_reducing_shared_output_twice(): + class FixedRouter(torch.nn.Module): + def forward(self, hidden_states): + tokens = hidden_states.shape[0] + return ( + torch.zeros(tokens, 4), + torch.full((tokens, 2), 0.5), + torch.zeros(tokens, 2, dtype=torch.int32), + ) + + class FixedExperts(torch.nn.Module): + def forward(self, hidden_states, _topk_ids, _topk_weights): + return torch.full_like(hidden_states, 2.0) + + class ZeroGate(torch.nn.Module): + def forward(self, hidden_states): + return torch.zeros(hidden_states.shape[0], 1) + + block = Qwen35MoeSparseMoeBlock.__new__(Qwen35MoeSparseMoeBlock) + torch.nn.Module.__init__(block) + block.shared_expert = torch.nn.Identity() + block.shared_expert_gate = ZeroGate() + block.gate = FixedRouter() + block.experts = FixedExperts() + block.parallel_context = SimpleNamespace( + world_all_reduce=Mock(side_effect=lambda tensor: tensor * 3) + ) + hidden_states = torch.full((2, 3), 4.0) + + output, *_ = block._forward_chunk(hidden_states) + + torch.testing.assert_close(output, torch.full_like(hidden_states, 8.0)) + block.parallel_context.world_all_reduce.assert_called_once() + reduced_input = block.parallel_context.world_all_reduce.call_args.args[0] + torch.testing.assert_close(reduced_input, torch.full_like(hidden_states, 2.0)) diff --git a/tests/test_tp_rpc.py b/tests/test_tp_rpc.py index 50842041..141ea2c9 100644 --- a/tests/test_tp_rpc.py +++ b/tests/test_tp_rpc.py @@ -258,6 +258,10 @@ def test_moe_workspace_warmup_uses_failure_synchronized_world_rpc(): assert "warmup_moe_workspace" in TP_RPC_STATUS_SYNC_METHODS +def test_runtime_diagnostics_use_failure_synchronized_world_rpc(): + assert "runtime_diagnostic_status" in TP_RPC_STATUS_SYNC_METHODS + + def test_fake_prefill_warmup_uses_failure_synchronized_world_rpc(): assert "set_warmup_fake_prefill_attention" in TP_RPC_STATUS_SYNC_METHODS diff --git a/tests/test_triton_moe.py b/tests/test_triton_moe.py index acd90195..7e101e1f 100644 --- a/tests/test_triton_moe.py +++ b/tests/test_triton_moe.py @@ -93,14 +93,15 @@ def test_moe_alignment_covers_hotspot_and_empty_rank(dtype): assert int(valid.numel()) == expected +@pytest.mark.parametrize("num_experts", [128, 256]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("norm_topk_prob", [False, True]) @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for Triton MoE tests.") -def test_topk_softmax_matches_pytorch(dtype, norm_topk_prob): +def test_topk_softmax_matches_pytorch(num_experts, dtype, norm_topk_prob): torch.manual_seed(21) - base = torch.arange(128, dtype=dtype, device="cuda") / 16 - 4 + base = torch.arange(num_experts, dtype=dtype, device="cuda") / 16 - 4 logits = torch.stack( - [base[torch.randperm(128, device="cuda")] for _ in range(257)] + [base[torch.randperm(num_experts, device="cuda")] for _ in range(257)] ) expected_probs = torch.softmax(logits, dim=-1, dtype=torch.float32) expected_weights, expected_ids = torch.topk(expected_probs, 8, dim=-1) @@ -166,7 +167,7 @@ def test_topk_softmax_nonfinite_inputs_keep_ids_in_range_and_propagate_nan(): @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for Triton MoE tests.") def test_topk_softmax_rejects_unsupported_shape_and_layout(): - with pytest.raises(ValueError, match="num_experts=128"): + with pytest.raises(ValueError, match="num_experts in"): topk_softmax( torch.zeros(2, 64, dtype=torch.bfloat16, device="cuda"), top_k=8, From e3b702a6a40c1f3063a30754dc8ce9cd1cf503dc Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Mon, 10 Aug 2026 23:45:12 +0800 Subject: [PATCH 02/35] feat: support qwen3.6 moe fp8 --- README.md | 6 +- docs/en/features/supported-models.md | 15 +- docs/zh/features/supported-models.md | 14 +- .../validation/qwen36_compare_longbench.py | 52 +++- src/sparsevllm/configs/model.py | 33 ++- src/sparsevllm/models/qwen3_5_moe.py | 223 +++++++++++++++++- tests/test_qwen35_moe.py | 128 +++++++++- tests/test_qwen36_longbench_compare.py | 46 ++++ 8 files changed, 473 insertions(+), 44 deletions(-) create mode 100644 tests/test_qwen36_longbench_compare.py diff --git a/README.md b/README.md index 9350c932..0ed85ef7 100644 --- a/README.md +++ b/README.md @@ -69,7 +69,7 @@ Read the method overview and integration rules in | Qwen3 | ✅ | | Qwen3MoE | ✅ | | Qwen3.5 / Qwen3.6 | ✅ | -| Qwen3.6 MoE | ✅ (BF16 text-only) | +| Qwen3.6 MoE | ✅ (BF16 / block FP8, text-only) | | Llama 3 / 3.1 | ✅ | | MiniMax M2.7 | ✅ | @@ -159,8 +159,8 @@ 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. +Sparse-vLLM supports Qwen3.5/Qwen3.6 dense and MoE 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/docs/en/features/supported-models.md b/docs/en/features/supported-models.md index 748c0ffd..f4aa71eb 100644 --- a/docs/en/features/supported-models.md +++ b/docs/en/features/supported-models.md @@ -16,7 +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 only | ✅ | 1 only | ✅ | +| Qwen3.6 MoE | `qwen3_5_moe` | BF16 / block FP8 | 1/2 | 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 | ✅ | @@ -34,13 +34,16 @@ the existing EP layout uses world size `E`. Qwen3.6 MoE always uses the outer-TP layout: attention and Gated DeltaNet TP are `T`, MoE EP is `E`, MoE TP is `T / E`, and world size is `T`. It requires -`DP=1`, `T % E == 0`, and BF16 language-model weights. The first release is a -text-only CausalLM runtime, rejects image/video and MTP inputs, supports only -Vanilla KV runtime, and captures decode (not prefill) with CUDA Graph. +`DP=1`, `T % E == 0`, and BF16 activations with either BF16 or block FP8 +language-model weights. The runtime is text-only, rejects image/video and MTP +inputs, supports only Vanilla KV runtime, and captures decode (not prefill) +with CUDA Graph. Outer TP is limited to 1 or 2 by the two KV heads; FP8 also +requires every TP-local quantized Linear dimension to remain 128-aligned. 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 diff --git a/docs/zh/features/supported-models.md b/docs/zh/features/supported-models.md index eab75258..fc5e13db 100644 --- a/docs/zh/features/supported-models.md +++ b/docs/zh/features/supported-models.md @@ -12,7 +12,7 @@ | 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 | ✅ | 仅支持 1 | ✅ | +| Qwen3.6 MoE | `qwen3_5_moe` | BF16 / 块级 FP8 | 1/2 | 仅支持 1 | ✅ | | Llama 3 / 3.1 | `llama` | BF16 / FP16 | ✅ | 仅支持 1 | 仅支持 1 | | MiniMax M2.7 | `minimax_m2` | 块级 FP8，非量化权重使用 BF16 | ✅ | 仅支持 1 | ✅ | @@ -27,11 +27,15 @@ TP 要求模型 dtype 为 BF16；FP16 Qwen3MoE checkpoint 仅支持 `TP=1`。当 Qwen3.6 MoE 始终使用 outer-TP 布局：attention 与 Gated DeltaNet TP 为 `T`、MoE EP 为 `E`、MoE TP 为 `T / E`，world size 为 `T`。该模型要求 -`DP=1`、`T % E == 0` 且语言模型权重为 BF16。首版仅支持纯文本 CausalLM， -明确拒绝 image/video 与 MTP 输入，只支持 Vanilla KV runtime；CUDA Graph -仅覆盖 decode，不覆盖 prefill。 +`DP=1`、`T % E == 0`，激活为 BF16，语言模型权重可使用 BF16 或块级 FP8。 +当前 runtime 仅支持纯文本 CausalLM，明确拒绝 image/video 与 MTP 输入，只 +支持 Vanilla KV runtime；CUDA Graph 仅覆盖 decode，不覆盖 prefill。两个 +KV heads 将 outer TP 限定为 1 或 2；FP8 还要求所有 TP-local 量化 Linear +维度保持 128 对齐。 -块级 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`。 ## 稀疏方法支持 diff --git a/scripts/validation/qwen36_compare_longbench.py b/scripts/validation/qwen36_compare_longbench.py index f7476826..3053ce89 100644 --- a/scripts/validation/qwen36_compare_longbench.py +++ b/scripts/validation/qwen36_compare_longbench.py @@ -30,7 +30,38 @@ def _read_jsonl(path: Path) -> list[dict[str, Any]]: return rows -def _load_run(path: Path) -> dict[str, Any]: +def _validate_worker_providers( + workers: list[dict[str, Any]], + *, + precision: str, + path: Path, +) -> None: + if precision == "bf16": + valid = all( + worker.get("moe_expert_provider") == "triton" + and worker.get("moe_router_provider") == "triton" + for worker in workers + ) + else: + valid_expert_providers = { + "flashinfer_cutlass_fp8_sm90", + "triton", + } + valid = all( + worker.get("moe_expert_provider") in valid_expert_providers + and worker.get("moe_router_provider") == "triton" + and worker.get("moe_weight_dtype") == "torch.float8_e4m3fn" + and isinstance(worker.get("fp8_linear_provider"), str) + and bool(worker["fp8_linear_provider"]) + for worker in workers + ) + if not valid: + raise RuntimeError( + f"LongBench run {path} has invalid {precision.upper()} providers." + ) + + +def _load_run(path: Path, *, precision: str) -> dict[str, Any]: required = ( "resolved_config.json", "raw_outputs.jsonl", @@ -57,12 +88,7 @@ def _load_run(path: Path) -> dict[str, Any]: workers = graph.get("workers") if not isinstance(workers, list) or not workers: raise RuntimeError(f"LongBench run {path} has no worker status records.") - if any( - worker.get("moe_expert_provider") != "triton" - or worker.get("moe_router_provider") != "triton" - for worker in workers - ): - raise RuntimeError(f"LongBench run {path} did not use Triton providers.") + _validate_worker_providers(workers, precision=precision, path=path) config = _read_json(path / "resolved_config.json") expected_per_task = int(config["selection"]["samples_per_task"]) counts = { @@ -136,12 +162,17 @@ def main() -> int: parser.add_argument("--single", type=Path, required=True) parser.add_argument("--pure-tp", type=Path, required=True) parser.add_argument("--tp-ep", type=Path, required=True) + parser.add_argument( + "--precision", + choices=("bf16", "fp8"), + default="bf16", + ) parser.add_argument("--output-dir", type=Path, required=True) args = parser.parse_args() runs = { - "single": _load_run(args.single.resolve()), - "pure_tp": _load_run(args.pure_tp.resolve()), - "tp_ep": _load_run(args.tp_ep.resolve()), + "single": _load_run(args.single.resolve(), precision=args.precision), + "pure_tp": _load_run(args.pure_tp.resolve(), precision=args.precision), + "tp_ep": _load_run(args.tp_ep.resolve(), precision=args.precision), } output_dir = args.output_dir.resolve() output_dir.mkdir(parents=True, exist_ok=False) @@ -252,6 +283,7 @@ def main() -> int: ["git", "rev-parse", "HEAD"], text=True, capture_output=True, check=True ).stdout.strip(), "sources": {topology: str(run["path"]) for topology, run in runs.items()}, + "precision": args.precision, } _write_json(output_dir / "run_info.json", run_info) _write_json(output_dir / "raw_outputs.json", run_info["sources"]) diff --git a/src/sparsevllm/configs/model.py b/src/sparsevllm/configs/model.py index 7c148d56..4b1b794c 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -250,16 +250,13 @@ def _validate_qwen35_moe_checkpoint_config( "architectures=['Qwen3_5MoeForConditionalGeneration'], " f"got {list(architectures)}." ) - if quantization_config.enabled: - raise NotImplementedError( - "Qwen3.6 MoE v1 supports BF16 checkpoints only; FP8 is out of scope." - ) 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 NotImplementedError( - "Qwen3.6 MoE v1 requires BF16 language-model weights, " + "Qwen3.6 MoE requires BF16 activations with either BF16 or block-FP8 " + "language-model weights, " f"got dtype={configured_dtype!r}." ) for field_name, expected in _QWEN35_MOE_FIXED_FIELDS.items(): @@ -825,9 +822,33 @@ def _validate_qwen35_moe_runtime(config, *, model_type: str) -> None: ) if getattr(hf_config, "torch_dtype", None) != torch.bfloat16: raise NotImplementedError( - "Qwen3.6 MoE v1 requires BF16 weights, got " + "Qwen3.6 MoE requires BF16 activations, got " f"torch_dtype={getattr(hf_config, 'torch_dtype', None)}." ) + if config.quantization_config.enabled: + block_size = tuple(config.quantization_config.weight_block_size or ()) + if block_size != (128, 128): + raise ValueError( + "Qwen3.6 MoE FP8 requires weight_block_size=(128, 128), " + f"got {block_size}." + ) + fp8_local_dimensions = { + "hidden_size": int(hf_config.hidden_size), + "shared_expert_intermediate_size": int( + hf_config.shared_expert_intermediate_size + ) + // outer_tp_size, + } + invalid_fp8_dimensions = { + name: value + for name, value in fp8_local_dimensions.items() + if value % 128 + } + if invalid_fp8_dimensions: + raise ValueError( + "Qwen3.6 MoE FP8 local Linear dimensions must be 128-aligned, " + f"got TP={outer_tp_size}, invalid={invalid_fp8_dimensions}." + ) _validate_runtime_compatibility(config, model_type=model_type) diff --git a/src/sparsevllm/models/qwen3_5_moe.py b/src/sparsevllm/models/qwen3_5_moe.py index ce0cc121..df995fe6 100644 --- a/src/sparsevllm/models/qwen3_5_moe.py +++ b/src/sparsevllm/models/qwen3_5_moe.py @@ -38,6 +38,14 @@ 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): @@ -72,28 +80,40 @@ def forward( class Qwen35MoePackedExperts(Qwen3MoePackedExperts): - """Qwen3.6 packed-3D checkpoint adapter for the shared MoE provider.""" + """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) - if self.fp8_enabled: - raise NotImplementedError( - "Qwen3.6 MoE v1 accepts BF16 expert weights only." - ) self._loaded_packed_projections: set[str] = set() def rank_local_weight_slice( self, source_shape: tuple[int, ...], *, - loaded_shard_id: str, + 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 = { @@ -129,6 +149,10 @@ def load_packed_expert_weight( 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( @@ -185,6 +209,9 @@ def load_packed_expert_weight( 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( @@ -319,6 +346,7 @@ class Qwen35MoeForCausalLM(Qwen35ForCausalLM): special_weight_loaders = { **Qwen35ForCausalLM.special_weight_loaders, ".packed_expert_weight": "load_packed_expert_weight", + ".expert_weight": "load_expert_weight", } def __init__(self, config) -> None: @@ -333,6 +361,8 @@ def __init__(self, config) -> None: 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: @@ -361,9 +391,16 @@ def recurrent_state_spec(config, attention_tp_size: int) -> RecurrentStateSpec: def runtime_diagnostic_status(self) -> dict[str, object]: experts = self.model.layers[0].mlp.experts router = self.model.layers[0].mlp.gate + shared_linear = self.model.layers[0].mlp.shared_expert.gate_up_proj return { "moe_expert_provider": experts.provider.name, "moe_router_provider": router.provider.name, + "moe_weight_dtype": str(experts.w13_weight.dtype), + "fp8_linear_provider": ( + shared_linear.quant_provider.name + if shared_linear.quantized + else None + ), "local_expert_start": int(experts.local_expert_start), "local_expert_end": int(experts.local_expert_end), } @@ -397,10 +434,32 @@ def warmup_moe(self, num_tokens: int = 1) -> None: block(hidden_states) device_runtime.synchronize() - def map_weight_name(self, source_weight_name: str) -> str: + 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" @@ -411,6 +470,13 @@ 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() @@ -426,6 +492,16 @@ def load_special_weight( 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: @@ -455,16 +531,95 @@ def record_skipped_weight( loaded_scale_shape: tuple[int, ...] | None, loaded_scale_dtype: str | None, ) -> None: - del loaded_weight_shape, loaded_weight_dtype + 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}." ) - if loaded_scale_shape is not None or loaded_scale_dtype is not None: + 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/MTP intentional skips must not consume " + "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: @@ -504,6 +659,49 @@ def validate_loaded_weights(self, loaded_parameter_names: set[str]) -> None: 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.")) @@ -520,13 +718,14 @@ def validate_loaded_weights(self, loaded_parameter_names: set[str]) -> None: "Qwen3.6 MoE checkpoint is missing expected intentional-skip " f"groups: {missing_skip_groups}." ) - first_experts = self.model.layers[0].mlp.experts logger.info( - "Loaded Qwen3.6 MoE rank {} expert_provider={} router_provider={} " + "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, diff --git a/tests/test_qwen35_moe.py b/tests/test_qwen35_moe.py index d370bd4a..21f46544 100644 --- a/tests/test_qwen35_moe.py +++ b/tests/test_qwen35_moe.py @@ -21,7 +21,7 @@ ) -def _outer_config(): +def _outer_config(*, fp8: bool = False): layer_types = [ "full_attention" if (layer_idx + 1) % 4 == 0 else "linear_attention" for layer_idx in range(40) @@ -55,11 +55,20 @@ def _outer_config(): torch_dtype=torch.bfloat16, quantization_config=None, ) - return SimpleNamespace( + outer_config = SimpleNamespace( model_type="qwen3_5_moe", architectures=["Qwen3_5MoeForConditionalGeneration"], text_config=text_config, ) + if fp8: + del text_config.quantization_config + outer_config.quantization_config = { + "quant_method": "fp8", + "fmt": "e4m3", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + } + return outer_config def _make_config(tmp_path, **overrides): @@ -104,6 +113,40 @@ def _pure_tp_context(world_rank: int) -> ParallelContext: ) +def _fp8_expert_config(): + return SimpleNamespace( + num_experts=2, + hidden_size=128, + moe_intermediate_size=128, + num_experts_per_tok=1, + dtype=torch.bfloat16, + torch_dtype=torch.bfloat16, + decode_cuda_graph=True, + quantization_config=QuantizationConfig( + enabled=True, + quant_method="fp8", + weight_dtype="e4m3", + activation_scheme="dynamic", + weight_block_size=(128, 128), + model_name="Qwen3.6 MoE", + ), + ) + + +def _make_fp8_experts(): + with ( + patch( + "sparsevllm.models.qwen3_moe.get_parallel_context", + return_value=_single_context(), + ), + patch( + "sparsevllm.models.qwen3_moe.resolve_moe_provider", + return_value=TritonMoeProvider(), + ), + ): + return Qwen35MoePackedExperts(_fp8_expert_config()) + + def test_qwen36_moe_config_normalizes_text_runtime_and_topology(tmp_path): config = _make_config( tmp_path, @@ -148,6 +191,27 @@ def test_qwen36_moe_rejects_non_bf16_checkpoint(tmp_path): Config(model=str(tmp_path)) +def test_qwen36_moe_accepts_outer_block_fp8_config(tmp_path): + with patch( + "sparsevllm.configs.runtime.AutoConfig.from_pretrained", + return_value=_outer_config(fp8=True), + ): + config = Config(model=str(tmp_path)) + + assert config.quantization_config.enabled is True + assert config.quantization_config.weight_dtype == "e4m3" + assert config.quantization_config.weight_block_size == (128, 128) + + +def test_qwen36_moe_fp8_rejects_unsupported_outer_tp(tmp_path): + with patch( + "sparsevllm.configs.runtime.AutoConfig.from_pretrained", + return_value=_outer_config(fp8=True), + ): + with pytest.raises(ValueError, match="num_key_value_heads"): + Config(model=str(tmp_path), tensor_parallel_size=8) + + def test_qwen36_moe_recurrent_state_uses_attention_tp_and_fp32_state(): config = _outer_config().text_config @@ -366,6 +430,66 @@ def test_packed_expert_pure_tp_shards_reconstruct_checkpoint(): torch.testing.assert_close(reconstructed_down, down) +def test_qwen36_fp8_experts_load_per_expert_weights_and_scales(): + experts = _make_fp8_experts() + sources = {} + for expert_id in range(experts.num_experts): + for projection in ("gate_proj", "up_proj", "down_proj"): + weight = ( + torch.randn(128, 128) + .clamp(-4.0, 4.0) + .to(torch.float8_e4m3fn) + ) + scale = torch.rand(1, 1, dtype=torch.bfloat16) + 0.1 + sources[(expert_id, projection)] = (weight, scale) + experts.load_expert_weight(expert_id, projection, weight, scale) + + experts.validate_loaded_weights() + for expert_id in range(experts.num_experts): + gate, gate_scale = sources[(expert_id, "gate_proj")] + up, up_scale = sources[(expert_id, "up_proj")] + down, down_scale = sources[(expert_id, "down_proj")] + assert torch.equal(experts.w13_weight[expert_id, :128], gate) + assert torch.equal(experts.w13_weight[expert_id, 128:], up) + assert torch.equal(experts.w13_scale_inv[expert_id, :1], gate_scale) + assert torch.equal(experts.w13_scale_inv[expert_id, 1:], up_scale) + assert torch.equal(experts.w2_weight[expert_id], down) + assert torch.equal(experts.w2_scale_inv[expert_id], down_scale) + + +def test_qwen36_checkpoint_adapter_keeps_fp8_layout_model_local(): + experts = _make_fp8_experts() + model = Qwen35MoeForCausalLM.__new__(Qwen35MoeForCausalLM) + torch.nn.Module.__init__(model) + model.model = torch.nn.Module() + model.model.layers = torch.nn.ModuleList([torch.nn.Module()]) + model.model.layers[0].mlp = torch.nn.Module() + model.model.layers[0].mlp.experts = experts + + source_name = ( + "model.language_model.layers.0.mlp.experts.1.gate_proj.weight" + ) + target_name = model.map_weight_name(source_name) + + assert target_name == "model.layers.0.mlp.experts.1.gate_proj.expert_weight" + target = model.resolve_special_weight(target_name) + assert target is not None + assert target.module is experts + assert target.shard_id == (1, "gate_proj") + with pytest.raises(ValueError, match="must use per-expert"): + model.map_weight_name( + "model.language_model.layers.0.mlp.experts.gate_up_proj" + ) + + +def test_qwen36_fp8_experts_require_scale(): + experts = _make_fp8_experts() + weight = torch.ones(128, 128).to(torch.float8_e4m3fn) + + with pytest.raises(ValueError, match="Missing FP8 weight_scale_inv"): + experts.load_expert_weight(0, "gate_proj", weight, None) + + def test_routed_output_reduces_without_reducing_shared_output_twice(): class FixedRouter(torch.nn.Module): def forward(self, hidden_states): diff --git a/tests/test_qwen36_longbench_compare.py b/tests/test_qwen36_longbench_compare.py new file mode 100644 index 00000000..50ad1263 --- /dev/null +++ b/tests/test_qwen36_longbench_compare.py @@ -0,0 +1,46 @@ +from pathlib import Path + +import pytest + +from scripts.validation.qwen36_compare_longbench import ( + _validate_worker_providers, +) + + +def test_fp8_longbench_accepts_registered_graph_providers(): + workers = [ + { + "moe_expert_provider": "flashinfer_cutlass_fp8_sm90", + "moe_router_provider": "triton", + "moe_weight_dtype": "torch.float8_e4m3fn", + "fp8_linear_provider": "flashinfer_sm90", + }, + { + "moe_expert_provider": "triton", + "moe_router_provider": "triton", + "moe_weight_dtype": "torch.float8_e4m3fn", + "fp8_linear_provider": "flashinfer_sm90", + }, + ] + + _validate_worker_providers( + workers, + precision="fp8", + path=Path("fp8-run"), + ) + + +def test_fp8_longbench_rejects_missing_fp8_diagnostics(): + workers = [ + { + "moe_expert_provider": "triton", + "moe_router_provider": "triton", + } + ] + + with pytest.raises(RuntimeError, match="invalid FP8 providers"): + _validate_worker_providers( + workers, + precision="fp8", + path=Path("fp8-run"), + ) From 5ea67dc054219353a053317d7fe724031f5bb859 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 00:44:39 +0800 Subject: [PATCH 03/35] feat: enable qwen3.6 moe sparse methods --- benchmark/sparsevllm_regression/manifest.json | 2 +- docs/en/features/supported-models.md | 9 ++-- docs/zh/features/supported-models.md | 10 ++--- src/sparsevllm/method_registry.py | 7 +++- tests/test_qwen35_moe.py | 42 +++++++++++++++++-- tests/test_sparsevllm_regression_grading.py | 10 ++++- 6 files changed, 64 insertions(+), 16 deletions(-) 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/docs/en/features/supported-models.md b/docs/en/features/supported-models.md index f4aa71eb..d9126a97 100644 --- a/docs/en/features/supported-models.md +++ b/docs/en/features/supported-models.md @@ -36,9 +36,10 @@ Qwen3.6 MoE always uses the outer-TP layout: attention and Gated DeltaNet TP are `T`, MoE EP is `E`, MoE TP is `T / E`, and world size is `T`. It requires `DP=1`, `T % E == 0`, and BF16 activations with either BF16 or block FP8 language-model weights. The runtime is text-only, rejects image/video and MTP -inputs, supports only Vanilla KV runtime, and captures decode (not prefill) -with CUDA Graph. Outer TP is limited to 1 or 2 by the two KV heads; FP8 also -requires every TP-local quantized Linear dimension to remain 128-aligned. +inputs, and captures decode (not prefill) with CUDA Graph. Sparse methods apply +only to the full-attention layers; Gated DeltaNet layers keep their recurrent +state path. Outer TP is limited to 1 or 2 by the two KV heads; FP8 also requires +every TP-local quantized Linear dimension to remain 128-aligned. Block FP8 support requires E4M3 weights, dynamic activation quantization, and a `128 x 128` weight block size. Qwen3.5/Qwen3.6 dense configurations are @@ -53,7 +54,7 @@ normalized internally to `model_type=qwen3_5`; Qwen3.6 MoE uses | Qwen3 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | Compressor required² | | Qwen3MoE | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Qwen3.5 / Qwen3.6 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | Matched checkpoint³ | -| Qwen3.6 MoE | ✅ | — | — | — | — | — | — | — | — | — | +| Qwen3.6 MoE | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Llama 3 / 3.1 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | Selected checkpoint¹ | Compressor required² | | MiniMax M2.7 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | diff --git a/docs/zh/features/supported-models.md b/docs/zh/features/supported-models.md index fc5e13db..5ba0f170 100644 --- a/docs/zh/features/supported-models.md +++ b/docs/zh/features/supported-models.md @@ -28,10 +28,10 @@ TP 要求模型 dtype 为 BF16；FP16 Qwen3MoE checkpoint 仅支持 `TP=1`。当 Qwen3.6 MoE 始终使用 outer-TP 布局：attention 与 Gated DeltaNet TP 为 `T`、MoE EP 为 `E`、MoE TP 为 `T / E`，world size 为 `T`。该模型要求 `DP=1`、`T % E == 0`，激活为 BF16，语言模型权重可使用 BF16 或块级 FP8。 -当前 runtime 仅支持纯文本 CausalLM，明确拒绝 image/video 与 MTP 输入，只 -支持 Vanilla KV runtime；CUDA Graph 仅覆盖 decode，不覆盖 prefill。两个 -KV heads 将 outer TP 限定为 1 或 2；FP8 还要求所有 TP-local 量化 Linear -维度保持 128 对齐。 +当前 runtime 仅支持纯文本 CausalLM，明确拒绝 image/video 与 MTP 输入。 +两个 KV heads 将 outer TP 限定为 1 或 2；CUDA Graph 仅覆盖 decode，不覆盖 +prefill。稀疏方法只作用于 full-attention 层，Gated DeltaNet 层继续使用独立 +的递归状态路径；FP8 还要求所有 TP-local 量化 Linear 维度保持 128 对齐。 块级 FP8 要求使用 E4M3 权重、动态激活量化以及 `128 x 128` 的权重块大小。 Qwen3.5/Qwen3.6 Dense 配置在内部统一规范为 `model_type=qwen3_5`；Qwen3.6 @@ -45,7 +45,7 @@ MoE 使用 `model_type=qwen3_5_moe`。 | Qwen3 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | 需要压缩器² | | Qwen3MoE | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Qwen3.5 / Qwen3.6 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | 匹配的 checkpoint³ | -| Qwen3.6 MoE | ✅ | — | — | — | — | — | — | — | — | — | +| Qwen3.6 MoE | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Llama 3 / 3.1 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | 指定 checkpoint¹ | 需要 compressor² | | MiniMax M2.7 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | diff --git a/src/sparsevllm/method_registry.py b/src/sparsevllm/method_registry.py index 2383b541..d95b0af8 100644 --- a/src/sparsevllm/method_registry.py +++ b/src/sparsevllm/method_registry.py @@ -62,6 +62,7 @@ "qwen3", "qwen3_moe", "qwen3_5", + "qwen3_5_moe", "llama", "minimax_m2", } @@ -143,10 +144,12 @@ class ModelRuntimeCompatibility: QWEN35_MOE_COMPATIBILITY = ModelRuntimeCompatibility( parallel_mode="outer_tp_moe_tp_ep", - sparse_methods=frozenset({""}), + sparse_methods=QWEN3_MOE_TP_EP_COMPATIBILITY.sparse_methods, prefix_cache_methods=frozenset(), requires_eager=False, - decode_cuda_graph_methods=frozenset({""}), + decode_cuda_graph_methods=( + QWEN3_MOE_TP_EP_COMPATIBILITY.decode_cuda_graph_methods + ), ) MINIMAX_M2_EP_COMPATIBILITY = ModelRuntimeCompatibility( diff --git a/tests/test_qwen35_moe.py b/tests/test_qwen35_moe.py index 21f46544..46ed1494 100644 --- a/tests/test_qwen35_moe.py +++ b/tests/test_qwen35_moe.py @@ -166,9 +166,45 @@ def test_qwen36_moe_config_normalizes_text_runtime_and_topology(tmp_path): assert config.runtime_layout.num_kv_layers == 10 -def test_qwen36_moe_rejects_non_vanilla_sparse_method(tmp_path): - with pytest.raises(ValueError, match="validated methods: 'vanilla'"): - _make_config(tmp_path, vllm_sparse_method="quest") +@pytest.mark.parametrize( + "method", + [ + "streamingllm", + "snapkv", + "h2o", + "pyramidkv", + "omnikv", + "quest", + "rkv", + ], +) +def test_qwen36_moe_accepts_asset_free_sparse_graph_methods(tmp_path, method): + config = _make_config( + tmp_path, + vllm_sparse_method=method, + full_attn_layers="3,11,19,27,35", + decode_cuda_graph=True, + enforce_eager=False, + ) + + assert config.vllm_sparse_method == method + assert config.decode_cuda_graph is True + + +@pytest.mark.parametrize( + ("method", "error"), + [ + ("skipkv", "official models with released steering vectors"), + ("deltakv", "validated methods"), + ], +) +def test_qwen36_moe_rejects_sparse_methods_requiring_model_assets( + tmp_path, + method, + error, +): + with pytest.raises(ValueError, match=error): + _make_config(tmp_path, vllm_sparse_method=method) def test_qwen36_moe_rejects_invalid_outer_tp_ep_topology(tmp_path): diff --git a/tests/test_sparsevllm_regression_grading.py b/tests/test_sparsevllm_regression_grading.py index c3dc692e..2a80f8a4 100644 --- a/tests/test_sparsevllm_regression_grading.py +++ b/tests/test_sparsevllm_regression_grading.py @@ -191,7 +191,15 @@ 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"]), From 6535f2b32c3310d69bcdc1184a7adba283f9c13c Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 01:04:54 +0800 Subject: [PATCH 04/35] refactor: remove qwen development code --- benchmark/long_bench/eval.py | 2 - benchmark/long_bench/pred.py | 173 +----- benchmark/microbench.py | 85 --- benchmark/runtime_validation.py | 25 - .../validation/qwen36_compare_artifacts.py | 379 ------------ .../validation/qwen36_compare_longbench.py | 302 ---------- .../validation/qwen36_end_to_end_reference.py | 526 ---------------- .../validation/qwen36_moe_bf16_reference.py | 367 ------------ .../validation/qwen36_summarize_microbench.py | 260 -------- .../run_qwen36_microbench_matrix.py | 391 ------------ src/sparsevllm/engine/llm_engine.py | 9 +- src/sparsevllm/engine/model_runner.py | 89 +-- src/sparsevllm/models/qwen3_5_moe.py | 98 +-- src/sparsevllm/operators/moe.py | 85 +-- src/sparsevllm/operators/moe_router.py | 52 +- src/sparsevllm/triton_kernel/moe_topk.py | 4 +- tests/test_input_processor.py | 54 -- tests/test_longbench_deltakv_contracts.py | 134 +---- tests/test_microbench_artifacts.py | 22 - tests/test_operator_providers.py | 119 +--- tests/test_prefill_schedule_policy.py | 6 +- tests/test_qwen35_mixed_runtime.py | 31 - tests/test_qwen35_moe.py | 563 ------------------ tests/test_qwen36_longbench_compare.py | 46 -- tests/test_sparsevllm_regression_grading.py | 10 +- tests/test_tp_rpc.py | 4 - tests/test_triton_moe.py | 9 +- 27 files changed, 74 insertions(+), 3771 deletions(-) delete mode 100644 benchmark/runtime_validation.py delete mode 100644 scripts/validation/qwen36_compare_artifacts.py delete mode 100644 scripts/validation/qwen36_compare_longbench.py delete mode 100644 scripts/validation/qwen36_end_to_end_reference.py delete mode 100644 scripts/validation/qwen36_moe_bf16_reference.py delete mode 100644 scripts/validation/qwen36_summarize_microbench.py delete mode 100644 scripts/validation/run_qwen36_microbench_matrix.py delete mode 100644 tests/test_input_processor.py delete mode 100644 tests/test_qwen35_moe.py delete mode 100644 tests/test_qwen36_longbench_compare.py diff --git a/benchmark/long_bench/eval.py b/benchmark/long_bench/eval.py index 1126ae71..e4698f02 100644 --- a/benchmark/long_bench/eval.py +++ b/benchmark/long_bench/eval.py @@ -327,7 +327,5 @@ def aggregate_category_scores(task_scores): json.dump(scores, f, ensure_ascii=False, indent=4) with open(os.path.join(path, "metrics.json"), "w") as f: json.dump(scores, f, ensure_ascii=False, indent=4) - with open(os.path.join(path, "aggregate_metrics.json"), "w") as f: - json.dump(scores, f, ensure_ascii=False, indent=4) if failed_tasks: raise SystemExit(1) diff --git a/benchmark/long_bench/pred.py b/benchmark/long_bench/pred.py index 9e3714ba..59fc46fe 100644 --- a/benchmark/long_bench/pred.py +++ b/benchmark/long_bench/pred.py @@ -1,6 +1,5 @@ import os import json -import hashlib import sys import subprocess import re @@ -24,7 +23,6 @@ from transformers import AutoTokenizer, GenerationConfig import torch.distributed as dist from benchmark.model_adapters.sparsevllm import get_sparsevllm_generate_api -from benchmark.runtime_validation import collect_worker_runtime_status from datetime import datetime BASE_PATH = os.getenv("SPARSEVLLM_OUTPUT_DIR", str(REPO_ROOT / "outputs")) @@ -40,30 +38,6 @@ } -def _sha256(path: str | os.PathLike[str]) -> str: - digest = hashlib.sha256() - with open(path, "rb") as handle: - for chunk in iter(lambda: handle.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _sha256_text(value: str) -> str: - return hashlib.sha256(value.encode("utf-8")).hexdigest() - - -def _git_value(*args: str) -> str | None: - result = subprocess.run( - ["git", *args], - cwd=REPO_ROOT, - text=True, - capture_output=True, - check=False, - ) - value = result.stdout.strip() - return value or None - - def get_longbench_data_path(dataset, use_longbench_e): if not DATA_PREFIX_PATH: raise FileNotFoundError( @@ -171,7 +145,6 @@ def _artifact_paths(out_root: str) -> dict[str, str]: "raw": os.path.join(out_root, "raw_outputs.jsonl"), "parsed": os.path.join(out_root, "parsed_outputs.jsonl"), "sample": os.path.join(out_root, "sample_results.jsonl"), - "per_sample": os.path.join(out_root, "per_sample_results.jsonl"), } @@ -188,23 +161,28 @@ def _write_decode_cuda_graph_status( "cannot verify decode CUDA graph execution." ) - statuses = collect_worker_runtime_status(llm) - if not statuses: - raise RuntimeError("Runtime validation returned no worker status records.") - configured_flags = [ - bool(status.get("decode_cuda_graph_configured")) for status in statuses - ] - configured = all(configured_flags) - expected = bool(getattr(llm.config, "decode_cuda_graph", False)) - all_active = all( - bool(status.get("decode_cuda_graph_active")) for status in statuses + runner = getattr(llm, "model_runner", None) + graph_runner = getattr(runner, "decode_cuda_graph_runner", None) + graph_states = ( + getattr(graph_runner, "_graphs", {}) + if graph_runner is not None + else {} + ) + graph_count = sum( + getattr(state, "graph", None) is not None + for state in graph_states.values() ) graph_status = { - "launcher_rank": int(rank), - "expected": expected, - "configured_on_all_workers": configured, - "active_on_all_workers": all_active, - "workers": statuses, + "rank": int(rank), + "configured": bool( + getattr(getattr(llm, "config", None), "decode_cuda_graph", False) + ), + "runner_initialized": graph_runner is not None, + "state_count": int(len(graph_states)), + "graph_count": int(graph_count), + "active": bool(graph_count > 0), + "last_state_key": str(getattr(graph_runner, "last_state_key", None)), + "state_keys": [str(key) for key in graph_states], } status_path = os.path.join( out_root, @@ -213,16 +191,6 @@ def _write_decode_cuda_graph_status( with open(status_path, "w", encoding="utf-8") as handle: json.dump(graph_status, handle, ensure_ascii=False, indent=2) handle.write("\n") - if expected and not configured: - raise RuntimeError( - "Decode CUDA Graph was requested but is not configured on every worker: " - f"{statuses!r}." - ) - if configured and not all_active: - raise RuntimeError( - "Decode CUDA Graph was requested but was not active on every model worker: " - f"{statuses!r}." - ) return graph_status @@ -266,8 +234,6 @@ def _write_sample_record( "source_idx", "status", "prompt_tokens", - "rendered_prompt", - "rendered_prompt_sha256", "raw_pred", "error", "traceback", @@ -290,7 +256,6 @@ def _write_sample_record( _append_jsonl(paths["raw"], raw_record) _append_jsonl(paths["parsed"], parsed_record) _append_jsonl(paths["sample"], record) - _append_jsonl(paths["per_sample"], record) # Keep the historical per-task files for benchmark/long_bench/eval.py. task_record = { @@ -409,10 +374,6 @@ def get_pred(rank, data, dataset_info, args, model, tokenizer, model_max_length, prompt_tokens=prompt_tokens, ) ) - prepared_records[-1]["rendered_prompt"] = prompt - prepared_records[-1]["rendered_prompt_sha256"] = _sha256_text( - prompt - ) except Exception as exc: record = _sample_base_record( dataset=dataset, @@ -530,28 +491,8 @@ def get_pred(rank, data, dataset_info, args, model, tokenizer, model_max_length, def worker(rank, world_size, datasets, dataset2prompt, dataset2maxlen, args, out_root, max_length_limit): - seed_everything(args.seed) + seed_everything(42) model, tokenizer, model_max_length, eos_token_ids = load_model_and_tokenizer(rank, args) - if rank == 0: - tokenizer_runtime = { - "tokenizer_class": type(tokenizer).__name__, - "tokenizer_path": args.tokenizer_path or args.model_path, - "chat_template": tokenizer.chat_template, - "bos_token": tokenizer.bos_token, - "bos_token_id": tokenizer.bos_token_id, - "eos_token": tokenizer.eos_token, - "eos_token_id": tokenizer.eos_token_id, - "effective_eos_token_ids": eos_token_ids, - "no_chat_template": args.no_chat_template, - "thinking_mode": args.thinking_mode, - } - with open( - os.path.join(out_root, "tokenizer_runtime.json"), - "w", - encoding="utf-8", - ) as handle: - json.dump(tokenizer_runtime, handle, ensure_ascii=False, indent=2) - handle.write("\n") for dataset in datasets: data_path = get_longbench_data_path(dataset, args.e) @@ -724,7 +665,6 @@ def parse_args(): parser.add_argument("--min_prompt_tokens", type=int, default=None) parser.add_argument("--samples_per_task", type=int, default=20) parser.add_argument("--min_required_samples", type=int, default=5) - parser.add_argument("--seed", type=int, default=20260810) parser.add_argument("--worker_rank", type=int, default=-1) parser.add_argument("--worker_world_size", type=int, default=1) parser.add_argument("--output_root", type=str, default=None) @@ -767,13 +707,7 @@ def parse_args(): for dataset in datasets: with open(os.path.join(out_root, f"{dataset}.jsonl"), 'w') as f: pass - for artifact in ( - "raw_outputs.jsonl", - "parsed_outputs.jsonl", - "sample_results.jsonl", - "per_sample_results.jsonl", - "longbench_mini_selection.jsonl", - ): + for artifact in ("raw_outputs.jsonl", "parsed_outputs.jsonl", "sample_results.jsonl", "longbench_mini_selection.jsonl"): with open(os.path.join(out_root, artifact), "w", encoding="utf-8") as f: pass @@ -781,79 +715,16 @@ def parse_args(): args.max_model_len = max_length_limit if args.worker_rank < 0: - dataset_files = [] - for dataset in datasets: - data_path = Path(get_longbench_data_path(dataset, args.e)).resolve() - dataset_files.append( - { - "dataset": dataset, - "path": str(data_path), - "size_bytes": data_path.stat().st_size, - "sha256": _sha256(data_path), - } - ) - model_path = Path(args.model_path).resolve() - model_files = {} - for name in ("config.json", "model.safetensors.index.json"): - path = model_path / name - if not path.is_file(): - raise FileNotFoundError(f"Required model metadata file is missing: {path}") - model_files[name] = {"path": str(path), "sha256": _sha256(path)} - tokenizer_files = {} - for name in ( - "tokenizer.json", - "tokenizer_config.json", - "generation_config.json", - "chat_template.jinja", - ): - path = model_path / name - if path.is_file(): - tokenizer_files[name] = { - "path": str(path), - "sha256": _sha256(path), - } - prompt_config_path = REPO_ROOT / "benchmark/long_bench/config/dataset2prompt.json" - maxlen_config_path = REPO_ROOT / "benchmark/long_bench/config/dataset2maxlen.json" resolved_config = { - "created_at": datetime.now().isoformat(timespec="seconds"), - "command": " ".join([sys.executable, *sys.argv]), - "git_commit": _git_value("rev-parse", "HEAD"), - "git_branch": _git_value("branch", "--show-current"), - "git_dirty": bool(_git_value("status", "--porcelain")), "model": args.model, "model_path": args.model_path, - "model_files": model_files, - "tokenizer_files": tokenizer_files, "tokenizer_path": args.tokenizer_path or args.model_path, "backend": "sparsevllm", - "provider_env": { - key: os.environ.get(key, "auto") - for key in ( - "SPARSEVLLM_MOE_PROVIDER", - "SPARSEVLLM_MOE_ROUTER_PROVIDER", - ) - }, "sparse_method": args.sparse_method, "deltakv_checkpoint_path": args.deltakv_checkpoint_path, "datasets": datasets, "longbench_data_root": DATA_PREFIX_PATH, - "dataset_files": dataset_files, - "prompt_config": { - "path": str(prompt_config_path), - "sha256": _sha256(prompt_config_path), - "selected_formats": { - dataset: dataset2prompt[dataset] for dataset in datasets - }, - }, - "maxlen_config": { - "path": str(maxlen_config_path), - "sha256": _sha256(maxlen_config_path), - "selected_values": { - dataset: dataset2maxlen[dataset] for dataset in datasets - }, - }, "max_model_len": args.max_model_len, - "seed": args.seed, "decoding": { "temperature": args.temperature, "top_p": args.top_p, diff --git a/benchmark/microbench.py b/benchmark/microbench.py index 610df79e..ffa5bf76 100644 --- a/benchmark/microbench.py +++ b/benchmark/microbench.py @@ -27,7 +27,6 @@ is_tp_decode_cuda_graph_supported, normalize_sparse_method, ) -from benchmark.runtime_validation import collect_worker_runtime_status DEFAULT_ALL_CHUNKED_PREFILL_SIZE = 96 * 1024 @@ -157,8 +156,6 @@ def _selected_env_snapshot() -> dict[str, str]: "SPARSEVLLM_LONG_PREFILL_OFFLOAD_MIN_TOKENS", "SPARSEVLLM_RAWKV_BUFFER_MODE", "SPARSEVLLM_RAWKV_PREFETCH", - "SPARSEVLLM_MOE_ROUTER_PROVIDER", - "SPARSEVLLM_MOE_PROVIDER", ] return {key: os.environ[key] for key in keys if key in os.environ} @@ -279,10 +276,6 @@ def _artifact_records(args, rows: list[dict[str, Any]]) -> list[dict[str, Any]]: "synchronize_step_timing", bool(getattr(args, "synchronize_step_timing", False)), ) - record.setdefault( - "warmup_output_len", - int(getattr(args, "warmup_output_len", 0) or 0), - ) if "prefill_tp" in row: record.setdefault("prefill_tok_s", row["prefill_tp"]) if "decode_tp" in row: @@ -293,8 +286,6 @@ def _artifact_records(args, rows: list[dict[str, Any]]) -> list[dict[str, Any]]: record.setdefault("itl_ms", row["itl"]) if "mem" in row: record.setdefault("peak_memory_gb", row["mem"]) - if "end_to_end_tp" in row: - record.setdefault("end_to_end_tok_s", row["end_to_end_tp"]) records.append(record) return records @@ -320,11 +311,9 @@ def _write_output_dir(args, rows: list[dict[str, Any]]) -> None: "output_len": int(args.output_len), "temperature": float(args.temperature), "top_p": float(args.top_p), - "seed": int(getattr(args, "seed", 20260810)), "synchronize_step_timing": bool( getattr(args, "synchronize_step_timing", False) ), - "warmup_output_len": int(getattr(args, "warmup_output_len", 0) or 0), "hyper_params": args.hyper_params_dict, "env": _selected_env_snapshot(), } @@ -402,10 +391,6 @@ def _decode_cuda_graph_status(llm) -> dict[str, Any]: } -def _worker_runtime_status(llm) -> list[dict[str, Any]]: - return collect_worker_runtime_status(llm) - - def _jsonable_config_value(value: Any) -> Any: if value is None or isinstance(value, (str, int, float, bool)): return value @@ -487,9 +472,6 @@ def _finished_outputs_have_tokens(finished_outputs) -> bool: def benchmark_task(method, length, bs, args, results_dict): - seed = int(getattr(args, "seed", 20260810)) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) # 为每个子进程重置显存统计 torch.cuda.reset_peak_memory_stats() torch.cuda.empty_cache() @@ -559,41 +541,6 @@ def benchmark_task(method, length, bs, args, results_dict): } llm = LLM(args.model_path, **engine_kwargs) resolved_engine_config = _resolved_engine_config(llm) - warmup_output_len = int(getattr(args, "warmup_output_len", 0) or 0) - if warmup_output_len < 0: - raise ValueError("warmup_output_len must be non-negative.") - if warmup_output_len: - warmup_prompt_len = min(int(length), 1024) - warmup_prompts = [[100] * warmup_prompt_len for _ in range(bs)] - warmup_sampling = [ - SamplingParams( - temperature=0.0, - top_p=1.0, - ignore_eos=True, - max_tokens=warmup_output_len, - ) - for _ in range(bs) - ] - warmup_outputs = llm.generate( - warmup_prompts, - warmup_sampling, - use_tqdm=False, - ) - if len(warmup_outputs) != bs or not llm.is_finished(): - raise RuntimeError( - "Microbenchmark warmup did not finish every request: " - f"expected={bs}, outputs={len(warmup_outputs)}, " - f"engine_finished={llm.is_finished()}." - ) - if bool(base_hyper_params.get("decode_cuda_graph")) and not all( - bool(status.get("decode_cuda_graph_active")) - for status in _worker_runtime_status(llm) - ): - raise RuntimeError( - "Decode CUDA Graph was requested but did not activate during warmup." - ) - torch.cuda.synchronize() - torch.cuda.reset_peak_memory_stats() prefix_cache_stats_before = _cache_stats(llm) prompt_token_ids = [[100] * length for _ in range(bs)] @@ -739,25 +686,8 @@ def add_wave(max_new_requests: int): t_end = perf_counter() duration = t_end - t_start - end_to_end_tokens = int(prefill_tokens + decode_tokens) - end_to_end_tp = end_to_end_tokens / duration if duration > 0 else 0.0 peak_mem = get_peak_memory() graph_status = _decode_cuda_graph_status(llm) - worker_runtime_status = _worker_runtime_status(llm) - primary_worker_status = worker_runtime_status[0] - provider_status = { - key: primary_worker_status[key] - for key in ("moe_expert_provider", "moe_router_provider") - if key in primary_worker_status - } - if bool(base_hyper_params.get("decode_cuda_graph")) and not all( - bool(status.get("decode_cuda_graph_active")) - for status in worker_runtime_status - ): - raise RuntimeError( - "Decode CUDA Graph was requested but was not active on every " - f"worker: {worker_runtime_status!r}." - ) prefix_cache_stats_after = _cache_stats(llm) prefix_cache_stats_delta = _numeric_delta(prefix_cache_stats_before, prefix_cache_stats_after) observed_prefix_hit_tokens = int(sum(prefix_hits_by_seq_id.values())) @@ -827,9 +757,6 @@ def add_wave(max_new_requests: int): "itl": avg_itl, "avg_bs": avg_active_bs, "mem": peak_mem, - "duration_s": duration, - "end_to_end_tokens": end_to_end_tokens, - "end_to_end_tp": end_to_end_tp, "has_queued": has_queued, "full_admission_reached": full_admission_reached, "impossible_full_admission": impossible_full_admission, @@ -846,8 +773,6 @@ def add_wave(max_new_requests: int): "scheduler_recompute_replays": recompute_replays, "decode_cuda_graph_expected": bool(base_hyper_params.get("decode_cuda_graph")), **graph_status, - **provider_status, - "worker_runtime_status": worker_runtime_status, "prefix_cache_required": bool(getattr(args, "require_prefix_cache_hit", False)), "prefix_cache_stats_before": prefix_cache_stats_before, "prefix_cache_stats_after": prefix_cache_stats_after, @@ -918,7 +843,6 @@ def main(): default=1.0, help="Nucleus sampling top-p. Only used when temperature > 0.", ) - parser.add_argument("--seed", type=int, default=20260810) parser.add_argument( "--admission_wave_size", type=int, @@ -945,15 +869,6 @@ def main(): "so post-sparse work is attributed to the step that launched it." ), ) - parser.add_argument( - "--warmup_output_len", - type=int, - default=0, - help=( - "Run an unmeasured same-batch warmup before each case; a positive value " - "also forces decode CUDA Graph capture before timed steps." - ), - ) parser.add_argument( "--wave_decode_gap_steps", type=int, diff --git a/benchmark/runtime_validation.py b/benchmark/runtime_validation.py deleted file mode 100644 index d6263017..00000000 --- a/benchmark/runtime_validation.py +++ /dev/null @@ -1,25 +0,0 @@ -from __future__ import annotations - -from typing import Any - - -def collect_worker_runtime_status(llm) -> list[dict[str, Any]]: - """Collect worker diagnostics without extending the public engine API.""" - model_runner = getattr(llm, "model_runner", None) - call = getattr(model_runner, "call", None) - if not callable(call): - raise RuntimeError( - "Sparse-VLLM runtime validation requires model_runner.call()." - ) - statuses = call("runtime_diagnostic_status") - expected = int(getattr(getattr(llm, "config", None), "world_size", 1)) - if ( - not isinstance(statuses, list) - or len(statuses) != expected - or not all(isinstance(status, dict) for status in statuses) - ): - raise RuntimeError( - "Runtime validation must return one status object per model worker: " - f"expected={expected}, got={statuses!r}." - ) - return statuses diff --git a/scripts/validation/qwen36_compare_artifacts.py b/scripts/validation/qwen36_compare_artifacts.py deleted file mode 100644 index 1d6c112d..00000000 --- a/scripts/validation/qwen36_compare_artifacts.py +++ /dev/null @@ -1,379 +0,0 @@ -from __future__ import annotations - -import argparse -import json -import subprocess -import sys -from datetime import datetime -from pathlib import Path -from typing import Any - -import torch -import torch.nn.functional as F - - -THRESHOLDS = { - "forced_logits_max_abs": 4.0, - "forced_logits_mean_abs": 0.6, - "near_tie_margin": 0.25, - "decoder_layer_cosine": 0.75, - "routing_layer_overlap": 0.90, - "routing_mean_overlap": 0.95, - "graph_eager_logits_max_abs": 0.0, -} - - -def _read_json(path: Path) -> dict[str, Any]: - value = json.loads(path.read_text(encoding="utf-8")) - if not isinstance(value, dict): - raise TypeError(f"Expected a JSON object in {path}, got {type(value).__name__}.") - return value - - -def _read_jsonl(path: Path) -> list[dict[str, Any]]: - rows = [] - for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): - try: - row = json.loads(line) - except json.JSONDecodeError as exc: - raise ValueError(f"Invalid JSONL at {path}:{line_no}: {exc}") from exc - if not isinstance(row, dict): - raise TypeError(f"Expected an object at {path}:{line_no}.") - rows.append(row) - return rows - - -def _load_tensor_list(path: Path) -> list[Any]: - value = torch.load(path, map_location="cpu", weights_only=False) - if not isinstance(value, list): - raise TypeError(f"Expected a tensor list in {path}, got {type(value).__name__}.") - return value - - -def _require_run(path: Path) -> None: - required = ( - "run_info.json", - "runtime_status.json", - "per_sample_results.jsonl", - "aggregate_metrics.json", - "raw_logits.pt", - ) - missing = [name for name in required if not (path / name).is_file()] - if missing: - raise FileNotFoundError(f"Validation run {path} is missing artifacts: {missing}.") - aggregate = _read_json(path / "aggregate_metrics.json") - if aggregate.get("status") != "success": - raise RuntimeError(f"Validation run {path} is not successful: {aggregate!r}.") - - -def _tensor_metrics(actual: torch.Tensor, reference: torch.Tensor) -> dict[str, float]: - if actual.shape != reference.shape: - raise ValueError( - f"Tensor shape mismatch: actual={tuple(actual.shape)}, reference={tuple(reference.shape)}." - ) - actual_fp32 = actual.float() - reference_fp32 = reference.float() - difference = (actual_fp32 - reference_fp32).abs() - return { - "max_abs": float(difference.max().item()), - "mean_abs": float(difference.mean().item()), - "cosine": float( - F.cosine_similarity(actual_fp32.flatten(), reference_fp32.flatten(), dim=0).item() - ), - } - - -def _output_ids(path: Path) -> list[list[int]]: - rows = _read_jsonl(path / "per_sample_results.jsonl") - if any(row.get("status") != "success" for row in rows): - raise RuntimeError(f"Non-success sample found in {path}.") - return [[int(token) for token in row["output_token_ids"]] for row in rows] - - -def _graph_records( - label: str, - eager_dir: Path, - graph_dir: Path, -) -> list[dict[str, Any]]: - eager_logits = _load_tensor_list(eager_dir / "raw_logits.pt") - graph_logits = _load_tensor_list(graph_dir / "raw_logits.pt") - if len(eager_logits) != len(graph_logits): - raise ValueError(f"{label} eager/Graph sample counts differ.") - max_abs = max( - _tensor_metrics(graph, eager)["max_abs"] - for eager, graph in zip(eager_logits, graph_logits) - ) - token_ids_equal = _output_ids(eager_dir) == _output_ids(graph_dir) - runtime = _read_json(graph_dir / "runtime_status.json") - workers = runtime.get("worker_runtime_status") - if not isinstance(workers, list) or not workers: - raise RuntimeError(f"{graph_dir} has no worker_runtime_status records.") - all_graph_active = all( - worker.get("decode_cuda_graph_configured") is True - and worker.get("decode_cuda_graph_active") is True - for worker in workers - ) - providers = { - (worker.get("moe_expert_provider"), worker.get("moe_router_provider")) - for worker in workers - } - status = ( - "success" - if max_abs <= THRESHOLDS["graph_eager_logits_max_abs"] - and token_ids_equal - and all_graph_active - and providers == {("triton", "triton")} - else "metric_failed" - ) - return [ - { - "check": "graph_eager_equivalence", - "topology": label, - "status": status, - "max_abs": max_abs, - "token_ids_equal": token_ids_equal, - "all_graph_active": all_graph_active, - "worker_count": len(workers), - "providers": sorted([list(pair) for pair in providers]), - } - ] - - -def _forced_reference_records( - label: str, - reference_dir: Path, - actual_dir: Path, -) -> list[dict[str, Any]]: - reference_logits = _load_tensor_list(reference_dir / "raw_logits.pt") - actual_logits = _load_tensor_list(actual_dir / "raw_logits.pt") - reference_hidden = _load_tensor_list( - reference_dir / "raw_cached_hidden_states.pt" - ) - actual_hidden = _load_tensor_list(actual_dir / "raw_cached_hidden_states.pt") - if not ( - len(reference_logits) - == len(actual_logits) - == len(reference_hidden) - == len(actual_hidden) - ): - raise ValueError(f"{label} forced-reference sample counts differ.") - - records: list[dict[str, Any]] = [] - for sample_idx, (reference, actual) in enumerate( - zip(reference_logits, actual_logits) - ): - metrics = _tensor_metrics(actual, reference) - reference_top2 = torch.topk(reference.float(), k=2) - actual_top1 = int(torch.argmax(actual).item()) - reference_top1 = int(reference_top2.indices[0].item()) - reference_margin = float( - (reference_top2.values[0] - reference_top2.values[1]).item() - ) - top1_acceptable = actual_top1 == reference_top1 or ( - reference_margin <= THRESHOLDS["near_tie_margin"] - and actual_top1 in {int(index) for index in reference_top2.indices.tolist()} - ) - status = ( - "success" - if metrics["max_abs"] <= THRESHOLDS["forced_logits_max_abs"] - and metrics["mean_abs"] <= THRESHOLDS["forced_logits_mean_abs"] - and top1_acceptable - else "metric_failed" - ) - records.append( - { - "check": "forced_prefix_logits", - "topology": label, - "sample_idx": sample_idx, - "status": status, - **metrics, - "reference_top1": reference_top1, - "actual_top1": actual_top1, - "reference_top1_margin": reference_margin, - "top1_acceptable": top1_acceptable, - } - ) - - reference_layers = reference_hidden[sample_idx] - actual_layers = actual_hidden[sample_idx] - if set(reference_layers) != set(actual_layers): - raise ValueError(f"{label} sample {sample_idx} hidden layer sets differ.") - for layer_idx in sorted(reference_layers): - layer_metrics = _tensor_metrics( - actual_layers[layer_idx], reference_layers[layer_idx] - ) - records.append( - { - "check": "decoder_layer_output", - "topology": label, - "sample_idx": sample_idx, - "layer_idx": int(layer_idx), - "layer_kind": ( - "embedding" - if layer_idx == -1 - else "final_norm" - if layer_idx == 40 - else "full_attention" - if layer_idx % 4 == 3 - else "gated_deltanet" - ), - "status": ( - "success" - if layer_metrics["cosine"] - >= THRESHOLDS["decoder_layer_cosine"] - else "metric_failed" - ), - **layer_metrics, - } - ) - return records - - -def _routing_records( - label: str, - reference_dir: Path, - actual_dir: Path, -) -> list[dict[str, Any]]: - reference = _load_tensor_list(reference_dir / "raw_moe_states.pt") - actual = _load_tensor_list(actual_dir / "raw_moe_states.pt") - if len(reference) != len(actual): - raise ValueError(f"{label} routing sample counts differ.") - records = [] - overlaps = [] - for sample_idx, (reference_layers, actual_layers) in enumerate( - zip(reference, actual) - ): - if set(reference_layers) != set(actual_layers): - raise ValueError(f"{label} sample {sample_idx} MoE layer sets differ.") - for layer_idx in sorted(reference_layers): - reference_ids = reference_layers[layer_idx]["topk_ids"] - actual_ids = actual_layers[layer_idx]["topk_ids"] - if reference_ids.shape != actual_ids.shape: - raise ValueError(f"{label} sample {sample_idx} layer {layer_idx} shape differs.") - row_overlaps = [ - len(set(left.tolist()) & set(right.tolist())) / reference_ids.shape[1] - for left, right in zip(reference_ids, actual_ids) - ] - overlap = float(sum(row_overlaps) / len(row_overlaps)) - overlaps.append(overlap) - records.append( - { - "check": "cross_topology_routing", - "topology": label, - "sample_idx": sample_idx, - "layer_idx": int(layer_idx), - "status": ( - "success" - if overlap >= THRESHOLDS["routing_layer_overlap"] - else "metric_failed" - ), - "topk_set_overlap": overlap, - "ordered_ids_equal": bool(torch.equal(reference_ids, actual_ids)), - } - ) - mean_overlap = float(sum(overlaps) / len(overlaps)) - records.append( - { - "check": "cross_topology_routing_aggregate", - "topology": label, - "status": ( - "success" - if mean_overlap >= THRESHOLDS["routing_mean_overlap"] - else "metric_failed" - ), - "mean_topk_set_overlap": mean_overlap, - "min_layer_topk_set_overlap": min(overlaps), - } - ) - return records - - -def _write_json(path: Path, value: Any) -> None: - path.write_text( - json.dumps(value, ensure_ascii=False, indent=2) + "\n", - encoding="utf-8", - ) - - -def main() -> None: - parser = argparse.ArgumentParser( - description="Compare fixed Qwen3.6 MoE correctness artifacts." - ) - parser.add_argument("--transformers", type=Path, required=True) - for name in ( - "single-eager", - "single-graph", - "tp-eager", - "tp-graph", - "tp-ep-eager", - "tp-ep-graph", - "forced-single", - "forced-tp", - "forced-tp-ep", - ): - parser.add_argument(f"--{name}", type=Path, required=True) - parser.add_argument("--output-dir", type=Path, required=True) - args = parser.parse_args() - - sources = { - key: Path(value).resolve() - for key, value in vars(args).items() - if key != "output_dir" - } - for path in sources.values(): - _require_run(path) - args.output_dir = args.output_dir.resolve() - args.output_dir.mkdir(parents=True, exist_ok=False) - - records = [] - records += _graph_records("single", sources["single_eager"], sources["single_graph"]) - records += _graph_records("pure_tp", sources["tp_eager"], sources["tp_graph"]) - records += _graph_records("tp_ep", sources["tp_ep_eager"], sources["tp_ep_graph"]) - for label, key in ( - ("single", "forced_single"), - ("pure_tp", "forced_tp"), - ("tp_ep", "forced_tp_ep"), - ): - records += _forced_reference_records( - label, sources["transformers"], sources[key] - ) - records += _routing_records( - "single_vs_pure_tp", sources["forced_single"], sources["forced_tp"] - ) - records += _routing_records( - "single_vs_tp_ep", sources["forced_single"], sources["forced_tp_ep"] - ) - - failed = [record for record in records if record["status"] != "success"] - aggregate = { - "status": "success" if not failed else "metric_failed", - "num_checks": len(records), - "success_checks": len(records) - len(failed), - "failed_checks": len(failed), - "thresholds": THRESHOLDS, - } - run_info = { - "created_at": datetime.now().isoformat(timespec="seconds"), - "command": " ".join(sys.argv), - "git_commit": subprocess.run( - ["git", "rev-parse", "HEAD"], text=True, capture_output=True, check=True - ).stdout.strip(), - "sources": {key: str(value) for key, value in sources.items()}, - "thresholds": THRESHOLDS, - } - _write_json(args.output_dir / "run_info.json", run_info) - _write_json(args.output_dir / "raw_outputs.json", run_info["sources"]) - _write_json(args.output_dir / "parsed_outputs.json", {"checks": records}) - with (args.output_dir / "per_sample_results.jsonl").open( - "w", encoding="utf-8" - ) as handle: - for record in records: - handle.write(json.dumps(record, ensure_ascii=False) + "\n") - _write_json(args.output_dir / "aggregate_metrics.json", aggregate) - print(json.dumps(aggregate, ensure_ascii=False, indent=2)) - if failed: - raise SystemExit(1) - - -if __name__ == "__main__": - main() diff --git a/scripts/validation/qwen36_compare_longbench.py b/scripts/validation/qwen36_compare_longbench.py deleted file mode 100644 index 3053ce89..00000000 --- a/scripts/validation/qwen36_compare_longbench.py +++ /dev/null @@ -1,302 +0,0 @@ -from __future__ import annotations - -import argparse -import json -import subprocess -import sys -from datetime import datetime -from pathlib import Path -from typing import Any - - -TASK_SCORE_MAX_DROP = 20.0 -MEAN_SCORE_MAX_DROP = 3.0 - - -def _read_json(path: Path) -> dict[str, Any]: - value = json.loads(path.read_text(encoding="utf-8")) - if not isinstance(value, dict): - raise TypeError(f"Expected a JSON object in {path}.") - return value - - -def _read_jsonl(path: Path) -> list[dict[str, Any]]: - rows = [] - for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): - value = json.loads(line) - if not isinstance(value, dict): - raise TypeError(f"Expected an object at {path}:{line_no}.") - rows.append(value) - return rows - - -def _validate_worker_providers( - workers: list[dict[str, Any]], - *, - precision: str, - path: Path, -) -> None: - if precision == "bf16": - valid = all( - worker.get("moe_expert_provider") == "triton" - and worker.get("moe_router_provider") == "triton" - for worker in workers - ) - else: - valid_expert_providers = { - "flashinfer_cutlass_fp8_sm90", - "triton", - } - valid = all( - worker.get("moe_expert_provider") in valid_expert_providers - and worker.get("moe_router_provider") == "triton" - and worker.get("moe_weight_dtype") == "torch.float8_e4m3fn" - and isinstance(worker.get("fp8_linear_provider"), str) - and bool(worker["fp8_linear_provider"]) - for worker in workers - ) - if not valid: - raise RuntimeError( - f"LongBench run {path} has invalid {precision.upper()} providers." - ) - - -def _load_run(path: Path, *, precision: str) -> dict[str, Any]: - required = ( - "resolved_config.json", - "raw_outputs.jsonl", - "parsed_outputs.jsonl", - "per_sample_results.jsonl", - "aggregate_metrics.json", - "decode_cuda_graph_status_rank0.json", - "tokenizer_runtime.json", - ) - missing = [name for name in required if not (path / name).is_file()] - if missing: - raise FileNotFoundError(f"LongBench run {path} is missing {missing}.") - metrics = _read_json(path / "aggregate_metrics.json") - if metrics.get("status") != "success": - raise RuntimeError(f"LongBench run {path} failed: {metrics!r}.") - samples = _read_jsonl(path / "per_sample_results.jsonl") - if not samples or any(sample.get("status") != "success" for sample in samples): - raise RuntimeError(f"LongBench run {path} contains non-success samples.") - graph = _read_json(path / "decode_cuda_graph_status_rank0.json") - if not graph.get("configured_on_all_workers") or not graph.get( - "active_on_all_workers" - ): - raise RuntimeError(f"LongBench run {path} did not activate Graph everywhere.") - workers = graph.get("workers") - if not isinstance(workers, list) or not workers: - raise RuntimeError(f"LongBench run {path} has no worker status records.") - _validate_worker_providers(workers, precision=precision, path=path) - config = _read_json(path / "resolved_config.json") - expected_per_task = int(config["selection"]["samples_per_task"]) - counts = { - str(dataset): sum( - sample.get("dataset") == dataset for sample in samples - ) - for dataset in config["datasets"] - } - if any(count != expected_per_task for count in counts.values()): - raise RuntimeError( - f"LongBench run {path} has incomplete task samples: " - f"expected_per_task={expected_per_task}, counts={counts}." - ) - return { - "path": path, - "config": config, - "metrics": metrics, - "samples": samples, - "graph": graph, - "tokenizer_runtime": _read_json(path / "tokenizer_runtime.json"), - } - - -def _selection(run: dict[str, Any]) -> list[tuple[str, int]]: - return [ - (str(sample["dataset"]), int(sample["source_idx"])) - for sample in run["samples"] - ] - - -def _rendered_prompt_hashes(run: dict[str, Any]) -> list[str]: - hashes = [sample.get("rendered_prompt_sha256") for sample in run["samples"]] - if any(not isinstance(value, str) or not value for value in hashes): - raise RuntimeError( - f"LongBench run {run['path']} is missing rendered prompt hashes." - ) - return hashes - - -def _dataset_fingerprints(run: dict[str, Any]) -> dict[str, str]: - return { - str(item["dataset"]): str(item["sha256"]) - for item in run["config"]["dataset_files"] - } - - -def _task_scores(run: dict[str, Any]) -> dict[str, float]: - datasets = [str(dataset) for dataset in run["config"]["datasets"]] - scores = {} - for dataset in datasets: - score = run["metrics"].get(dataset) - if not isinstance(score, (int, float)): - raise TypeError( - f"LongBench run {run['path']} has no numeric score for {dataset}: {score!r}." - ) - scores[dataset] = float(score) - return scores - - -def _write_json(path: Path, value: Any) -> None: - path.write_text( - json.dumps(value, ensure_ascii=False, indent=2) + "\n", - encoding="utf-8", - ) - - -def main() -> int: - parser = argparse.ArgumentParser( - description="Compare the fixed Qwen3.6 MoE LongBench subset across topologies." - ) - parser.add_argument("--single", type=Path, required=True) - parser.add_argument("--pure-tp", type=Path, required=True) - parser.add_argument("--tp-ep", type=Path, required=True) - parser.add_argument( - "--precision", - choices=("bf16", "fp8"), - default="bf16", - ) - parser.add_argument("--output-dir", type=Path, required=True) - args = parser.parse_args() - runs = { - "single": _load_run(args.single.resolve(), precision=args.precision), - "pure_tp": _load_run(args.pure_tp.resolve(), precision=args.precision), - "tp_ep": _load_run(args.tp_ep.resolve(), precision=args.precision), - } - output_dir = args.output_dir.resolve() - output_dir.mkdir(parents=True, exist_ok=False) - - reference = runs["single"] - reference_selection = _selection(reference) - reference_prompt_hashes = _rendered_prompt_hashes(reference) - reference_fingerprints = _dataset_fingerprints(reference) - reference_scores = _task_scores(reference) - reference_mean = sum(reference_scores.values()) / len(reference_scores) - records: list[dict[str, Any]] = [] - for topology, run in runs.items(): - config_equal = all( - run["config"].get(key) == reference["config"].get(key) - for key in ( - "datasets", - "seed", - "decoding", - "selection", - "model_files", - "tokenizer_files", - "provider_env", - "prompt_config", - "maxlen_config", - ) - ) - same_selection = _selection(run) == reference_selection - same_rendered_prompts = ( - _rendered_prompt_hashes(run) == reference_prompt_hashes - ) - same_fingerprints = _dataset_fingerprints(run) == reference_fingerprints - same_tokenizer_runtime = ( - run["tokenizer_runtime"] == reference["tokenizer_runtime"] - ) - records.append( - { - "check": "fixed_inputs", - "topology": topology, - "status": ( - "success" - if config_equal - and same_selection - and same_rendered_prompts - and same_fingerprints - and same_tokenizer_runtime - else "metric_failed" - ), - "config_equal": config_equal, - "same_sample_ids": same_selection, - "same_rendered_prompts": same_rendered_prompts, - "same_dataset_fingerprints": same_fingerprints, - "same_tokenizer_runtime": same_tokenizer_runtime, - "num_samples": len(run["samples"]), - } - ) - scores = _task_scores(run) - for task, reference_score in reference_scores.items(): - score = scores[task] - drop = reference_score - score - records.append( - { - "check": "task_quality", - "topology": topology, - "task": task, - "status": ( - "success" if drop <= TASK_SCORE_MAX_DROP else "metric_failed" - ), - "score": score, - "single_reference_score": reference_score, - "score_drop": drop, - "max_allowed_drop": TASK_SCORE_MAX_DROP, - } - ) - mean_score = sum(scores.values()) / len(scores) - mean_drop = reference_mean - mean_score - records.append( - { - "check": "mean_quality", - "topology": topology, - "status": ( - "success" if mean_drop <= MEAN_SCORE_MAX_DROP else "metric_failed" - ), - "mean_score": mean_score, - "single_reference_mean_score": reference_mean, - "score_drop": mean_drop, - "max_allowed_drop": MEAN_SCORE_MAX_DROP, - } - ) - - failed = [record for record in records if record["status"] != "success"] - aggregate = { - "status": "success" if not failed else "metric_failed", - "num_checks": len(records), - "success_checks": len(records) - len(failed), - "failed_checks": len(failed), - "thresholds": { - "task_score_max_drop": TASK_SCORE_MAX_DROP, - "mean_score_max_drop": MEAN_SCORE_MAX_DROP, - }, - "scores": { - topology: _task_scores(run) for topology, run in runs.items() - }, - } - run_info = { - "created_at": datetime.now().isoformat(timespec="seconds"), - "command": " ".join(sys.argv), - "git_commit": subprocess.run( - ["git", "rev-parse", "HEAD"], text=True, capture_output=True, check=True - ).stdout.strip(), - "sources": {topology: str(run["path"]) for topology, run in runs.items()}, - "precision": args.precision, - } - _write_json(output_dir / "run_info.json", run_info) - _write_json(output_dir / "raw_outputs.json", run_info["sources"]) - _write_json(output_dir / "parsed_outputs.json", {"checks": records}) - with (output_dir / "per_sample_results.jsonl").open( - "w", encoding="utf-8" - ) as handle: - for record in records: - handle.write(json.dumps(record, ensure_ascii=False) + "\n") - _write_json(output_dir / "aggregate_metrics.json", aggregate) - print(json.dumps(aggregate, ensure_ascii=False, indent=2)) - return 0 if not failed else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/validation/qwen36_end_to_end_reference.py b/scripts/validation/qwen36_end_to_end_reference.py deleted file mode 100644 index f62c2b8e..00000000 --- a/scripts/validation/qwen36_end_to_end_reference.py +++ /dev/null @@ -1,526 +0,0 @@ -from __future__ import annotations - -import argparse -import hashlib -import json -import os -import subprocess -import sys -import traceback -from datetime import datetime -from pathlib import Path -from typing import Any - -import torch -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( - Qwen3_5MoeRMSNormGated, - torch_causal_conv1d_update, - torch_chunk_gated_delta_rule, - torch_recurrent_gated_delta_rule, -) - - -REPO_ROOT = Path(__file__).resolve().parents[2] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -from benchmark.runtime_validation import collect_worker_runtime_status - -DEFAULT_PROMPTS = ( - "Sparse attention keeps the most useful context because", - "请用一句话解释专家并行的作用：", - "def fibonacci(n):\n \"\"\"Return the nth Fibonacci number.\"\"\"\n", -) - - -def _write_json(path: Path, value: Any) -> None: - path.write_text( - json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - - -def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: - with path.open("w", encoding="utf-8") as handle: - for row in rows: - handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n") - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - for chunk in iter(lambda: handle.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _git_value(*args: str) -> str | None: - result = subprocess.run( - ["git", *args], - cwd=REPO_ROOT, - check=False, - capture_output=True, - text=True, - ) - value = result.stdout.strip() - return value or None - - -def _tokenize_prompts(tokenizer, prompts: tuple[str, ...]) -> list[list[int]]: - token_ids = [] - for prompt in prompts: - add_special_tokens = True - if tokenizer.bos_token is None or prompt.startswith(tokenizer.bos_token): - add_special_tokens = False - encoded = tokenizer.encode(prompt, add_special_tokens=add_special_tokens) - if not encoded: - raise ValueError(f"Prompt tokenized to an empty sequence: {prompt!r}.") - token_ids.append([int(token_id) for token_id in encoded]) - return token_ids - - -def _run_transformers( - args, - tokenizer, - prompt_token_ids: list[list[int]], -) -> tuple[ - list[dict[str, Any]], - list[torch.Tensor], - dict[str, Any], - list[dict[int, torch.Tensor]], - list[dict[int, torch.Tensor]], - list[dict[int, dict[str, torch.Tensor]]], -]: - model = AutoModelForCausalLM.from_pretrained( - args.model, - dtype=torch.bfloat16, - local_files_only=True, - attn_implementation=("eager" if args.torch_reference_kernels else "sdpa"), - ).to("cuda").eval() - language_model = model.model - if hasattr(language_model, "language_model"): - language_model = language_model.language_model - if args.torch_reference_kernels: - for layer in language_model.layers: - linear_attn = getattr(layer, "linear_attn", None) - if linear_attn is None: - continue - linear_attn.causal_conv1d_fn = None - linear_attn.causal_conv1d_update = torch_causal_conv1d_update - linear_attn.chunk_gated_delta_rule = torch_chunk_gated_delta_rule - linear_attn.recurrent_gated_delta_rule = ( - torch_recurrent_gated_delta_rule - ) - torch_norm = Qwen3_5MoeRMSNormGated( - linear_attn.head_v_dim, - eps=linear_attn.layer_norm_epsilon, - ).to(device="cuda", dtype=torch.bfloat16) - with torch.no_grad(): - torch_norm.weight.copy_(linear_attn.norm.weight) - linear_attn.norm = torch_norm - rows = [] - logits = [] - hidden_snapshots: list[dict[int, torch.Tensor]] = [] - cached_hidden_snapshots: list[dict[int, torch.Tensor]] = [] - selected_layers = tuple(args.debug_hidden_layers) - live_hidden: dict[int, torch.Tensor] = {} - handles = [] - - def capture(layer_idx: int): - def hook(_module, _inputs, output): - tensor = output[0] if isinstance(output, tuple) else output - live_hidden[layer_idx] = tensor[:, -1].detach().cpu() - - return hook - - if selected_layers: - for layer_idx in selected_layers: - handles.append( - language_model.layers[layer_idx].register_forward_hook( - capture(layer_idx) - ) - ) - handles.append( - language_model.norm.register_forward_hook( - capture(len(language_model.layers)) - ) - ) - try: - for sample_id, (prompt, input_ids) in enumerate( - zip(DEFAULT_PROMPTS, prompt_token_ids) - ): - input_tensor = torch.tensor( - [input_ids], dtype=torch.long, device="cuda" - ) - with torch.inference_mode(): - generated = model.generate( - input_tensor, - do_sample=False, - max_new_tokens=args.max_new_tokens, - min_new_tokens=args.max_new_tokens, - use_cache=True, - pad_token_id=tokenizer.eos_token_id, - ) - output_ids = generated[0, input_tensor.shape[1] :].tolist() - if len(output_ids) != args.max_new_tokens: - raise RuntimeError( - f"Transformers sample {sample_id} generated {len(output_ids)} " - f"tokens, expected {args.max_new_tokens}." - ) - if selected_layers: - live_hidden[-1] = ( - language_model.embed_tokens(generated[:, -2]) - .detach() - .cpu() - ) - cached_hidden_snapshots.append( - dict(sorted(live_hidden.items())) - ) - final_input = generated[:, :-1] - live_hidden.clear() - final_logits = model( - input_ids=final_input, - use_cache=False, - return_dict=True, - ).logits[0, -1].detach().cpu() - if selected_layers: - live_hidden[-1] = ( - language_model.embed_tokens(final_input[:, -1]) - .detach() - .cpu() - ) - hidden_snapshots.append(dict(sorted(live_hidden.items()))) - logits.append(final_logits) - rows.append( - { - "sample_id": sample_id, - "status": "success", - "prompt": prompt, - "prompt_token_ids": input_ids, - "output_token_ids": [int(token_id) for token_id in output_ids], - "output_text": tokenizer.decode( - output_ids, skip_special_tokens=True - ), - } - ) - finally: - for handle in handles: - handle.remove() - return ( - rows, - logits, - {"backend": "transformers", "worker_runtime_status": []}, - hidden_snapshots, - cached_hidden_snapshots, - [], - ) - - -def _run_sparsevllm( - args, - tokenizer, - prompt_token_ids: list[list[int]], -) -> tuple[ - list[dict[str, Any]], - list[torch.Tensor], - dict[str, Any], - list[dict[int, torch.Tensor]], - list[dict[int, torch.Tensor]], - list[dict[int, dict[str, torch.Tensor]]], -]: - os.environ["SPARSEVLLM_DEBUG_RUNTIME"] = "1" - os.environ["SPARSEVLLM_DEBUG_MOE"] = "1" - if args.debug_hidden_layers: - os.environ["SPARSEVLLM_DEBUG_HIDDEN_LAYERS"] = ",".join( - str(layer_idx) for layer_idx in args.debug_hidden_layers - ) - if str(REPO_ROOT / "src") not in sys.path: - sys.path.insert(0, str(REPO_ROOT / "src")) - from sparsevllm import LLM, SamplingParams - - max_prompt_len = max(len(item) for item in prompt_token_ids) - llm = LLM( - model=str(args.model), - tensor_parallel_size=args.tensor_parallel_size, - expert_parallel_size=args.expert_parallel_size, - data_parallel_size=1, - enforce_eager=not args.decode_cuda_graph, - decode_cuda_graph=args.decode_cuda_graph, - gpu_memory_utilization=args.gpu_memory_utilization, - weight_loading_workers=args.weight_loading_workers, - max_model_len=max_prompt_len + args.max_new_tokens + 32, - max_num_seqs_in_batch=1, - max_decoding_seqs=1, - engine_prefill_chunk_size=max(64, max_prompt_len), - enable_profiler=False, - ) - rows = [] - logits = [] - debug_summaries = [] - hidden_snapshots = [] - moe_snapshots = [] - try: - sampling_params = SamplingParams( - temperature=0.0, - top_p=1.0, - max_tokens=args.max_new_tokens, - ignore_eos=True, - ) - for sample_id, (prompt, input_ids) in enumerate( - zip(DEFAULT_PROMPTS, prompt_token_ids) - ): - result = llm.generate( - [input_ids], sampling_params, use_tqdm=False - )[0] - output_ids = [int(token_id) for token_id in result["token_ids"]] - if len(output_ids) != args.max_new_tokens: - raise RuntimeError( - f"Sparse-vLLM sample {sample_id} generated {len(output_ids)} " - f"tokens, expected {args.max_new_tokens}." - ) - logits.append(llm.debug_last_logits().detach().cpu()[0]) - debug_summaries.append(llm.debug_sparse_state_summaries()) - if args.debug_hidden_layers: - hidden_snapshots.append(llm.debug_hidden_states()) - if args.debug_moe_states: - moe_snapshots.append(llm.debug_moe_states()) - rows.append( - { - "sample_id": sample_id, - "status": "success", - "prompt": prompt, - "prompt_token_ids": input_ids, - "output_token_ids": output_ids, - "output_text": result["text"], - } - ) - worker_status = collect_worker_runtime_status(llm) - if args.decode_cuda_graph and not all( - bool(status["decode_cuda_graph_active"]) for status in worker_status - ): - raise RuntimeError( - "Decode CUDA Graph is not active on every rank: " - f"{worker_status!r}." - ) - finally: - llm.exit() - return ( - rows, - logits, - { - "backend": "sparsevllm", - "worker_runtime_status": worker_status, - "debug_sparse_state_summaries": debug_summaries, - }, - hidden_snapshots, - hidden_snapshots, - moe_snapshots, - ) - - -def main() -> int: - parser = argparse.ArgumentParser( - description="Produce deterministic Qwen3.6 MoE end-to-end reference artifacts." - ) - parser.add_argument("--backend", choices=("transformers", "sparsevllm"), required=True) - parser.add_argument("--model", type=Path, required=True) - parser.add_argument("--output-dir", type=Path, required=True) - parser.add_argument("--max-new-tokens", type=int, default=8) - parser.add_argument("--seed", type=int, default=20260810) - parser.add_argument("--tensor-parallel-size", type=int, default=1) - parser.add_argument("--expert-parallel-size", type=int, default=1) - parser.add_argument("--decode-cuda-graph", action="store_true") - parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) - parser.add_argument("--weight-loading-workers", type=int, default=16) - parser.add_argument( - "--torch-reference-kernels", - action="store_true", - help=( - "For the Transformers backend, force eager full attention and the " - "explicit Torch Gated DeltaNet conv/chunk/recurrent/norm functions." - ), - ) - parser.add_argument( - "--debug-hidden-layers", - type=int, - nargs="*", - default=(), - help="Capture last-token hidden states after the selected decoder layers.", - ) - parser.add_argument( - "--debug-moe-states", - action="store_true", - help="Capture rank-0 per-layer MoE inputs, routing, and outputs.", - ) - parser.add_argument( - "--forced-prefix-artifact", - type=Path, - default=None, - help=( - "Append all but the last generated token from a prior " - "per_sample_results.jsonl artifact to each fixed prompt." - ), - ) - args = parser.parse_args() - if args.max_new_tokens <= 0: - raise ValueError("--max-new-tokens must be positive.") - if args.torch_reference_kernels and args.backend != "transformers": - raise ValueError("--torch-reference-kernels requires --backend transformers.") - invalid_hidden_layers = [ - layer_idx - for layer_idx in args.debug_hidden_layers - if layer_idx < 0 or layer_idx >= 40 - ] - if invalid_hidden_layers: - raise ValueError( - "--debug-hidden-layers must be in [0, 39], got " - f"{invalid_hidden_layers}." - ) - if not torch.cuda.is_available(): - raise RuntimeError("This validation requires CUDA.") - - args.model = args.model.resolve() - args.output_dir = args.output_dir.resolve() - if args.forced_prefix_artifact is not None: - args.forced_prefix_artifact = args.forced_prefix_artifact.resolve() - if not args.forced_prefix_artifact.is_file(): - raise FileNotFoundError( - "Forced-prefix artifact does not exist: " - f"{args.forced_prefix_artifact}." - ) - args.output_dir.mkdir(parents=True, exist_ok=False) - run_info = { - "created_at": datetime.now().isoformat(timespec="seconds"), - "command": " ".join(sys.argv), - "backend": args.backend, - "model": str(args.model), - "model_config_sha256": _sha256(args.model / "config.json"), - "model_index_sha256": _sha256( - args.model / "model.safetensors.index.json" - ), - "seed": args.seed, - "max_new_tokens": args.max_new_tokens, - "temperature": 0.0, - "top_p": 1.0, - "tensor_parallel_size": args.tensor_parallel_size, - "expert_parallel_size": args.expert_parallel_size, - "decode_cuda_graph": args.decode_cuda_graph, - "torch_reference_kernels": args.torch_reference_kernels, - "forced_prefix_artifact": ( - None - if args.forced_prefix_artifact is None - else str(args.forced_prefix_artifact) - ), - "forced_prefix_artifact_sha256": ( - None - if args.forced_prefix_artifact is None - else _sha256(args.forced_prefix_artifact) - ), - "gpu": torch.cuda.get_device_name(0), - "git_commit": _git_value("rev-parse", "HEAD"), - "git_branch": _git_value("branch", "--show-current"), - "git_dirty": bool(_git_value("status", "--porcelain")), - "requested_moe_provider": os.getenv("SPARSEVLLM_MOE_PROVIDER", "auto"), - "requested_moe_router_provider": os.getenv( - "SPARSEVLLM_MOE_ROUTER_PROVIDER", "auto" - ), - } - _write_json(args.output_dir / "run_info.json", run_info) - tokenizer = AutoTokenizer.from_pretrained(args.model, local_files_only=True) - prompt_token_ids = _tokenize_prompts(tokenizer, DEFAULT_PROMPTS) - if args.forced_prefix_artifact is not None: - forced_rows = [ - json.loads(line) - for line in args.forced_prefix_artifact.read_text( - encoding="utf-8" - ).splitlines() - if line.strip() - ] - if len(forced_rows) != len(prompt_token_ids): - raise ValueError( - "Forced-prefix artifact must contain one row per fixed prompt: " - f"expected={len(prompt_token_ids)} got={len(forced_rows)}." - ) - for sample_id, (token_ids, row) in enumerate( - zip(prompt_token_ids, forced_rows) - ): - output_token_ids = row.get("output_token_ids") - if not isinstance(output_token_ids, list) or len(output_token_ids) < 2: - raise ValueError( - f"Forced-prefix sample {sample_id} needs at least two " - "output_token_ids." - ) - token_ids.extend(int(token_id) for token_id in output_token_ids[:-1]) - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - - try: - if args.backend == "transformers": - ( - rows, - logits, - parsed, - hidden_snapshots, - cached_hidden_snapshots, - moe_snapshots, - ) = _run_transformers( - args, tokenizer, prompt_token_ids - ) - else: - ( - rows, - logits, - parsed, - hidden_snapshots, - cached_hidden_snapshots, - moe_snapshots, - ) = _run_sparsevllm( - args, tokenizer, prompt_token_ids - ) - except Exception as exc: - failure = { - "sample_id": None, - "status": "model_failed", - "error": repr(exc), - "traceback": traceback.format_exc(), - } - _write_jsonl(args.output_dir / "raw_outputs.jsonl", [failure]) - _write_jsonl(args.output_dir / "parsed_outputs.jsonl", [failure]) - _write_jsonl(args.output_dir / "per_sample_results.jsonl", [failure]) - _write_json(args.output_dir / "aggregate_metrics.json", failure) - raise - - _write_jsonl(args.output_dir / "raw_outputs.jsonl", rows) - _write_jsonl(args.output_dir / "parsed_outputs.jsonl", rows) - _write_jsonl(args.output_dir / "per_sample_results.jsonl", rows) - torch.save(logits, args.output_dir / "raw_logits.pt") - if hidden_snapshots: - torch.save( - hidden_snapshots, - args.output_dir / "raw_hidden_states.pt", - ) - if cached_hidden_snapshots: - torch.save( - cached_hidden_snapshots, - args.output_dir / "raw_cached_hidden_states.pt", - ) - if moe_snapshots: - torch.save( - moe_snapshots, - args.output_dir / "raw_moe_states.pt", - ) - _write_json(args.output_dir / "runtime_status.json", parsed) - aggregate = { - "status": "success", - "num_samples": len(rows), - "success_samples": sum(row["status"] == "success" for row in rows), - "failed_samples": sum(row["status"] != "success" for row in rows), - } - _write_json(args.output_dir / "aggregate_metrics.json", aggregate) - print(json.dumps(aggregate, ensure_ascii=False, indent=2)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/validation/qwen36_moe_bf16_reference.py b/scripts/validation/qwen36_moe_bf16_reference.py deleted file mode 100644 index a4db852d..00000000 --- a/scripts/validation/qwen36_moe_bf16_reference.py +++ /dev/null @@ -1,367 +0,0 @@ -from __future__ import annotations - -import argparse -import hashlib -import json -import os -import subprocess -import sys -from datetime import datetime -from pathlib import Path -from typing import Any - -import torch -import torch.nn.functional as F -from safetensors import safe_open -from transformers import AutoConfig -from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( - Qwen3_5MoeSparseMoeBlock, -) - - -REPO_ROOT = Path(__file__).resolve().parents[2] -if str(REPO_ROOT / "src") not in sys.path: - sys.path.insert(0, str(REPO_ROOT / "src")) - -from sparsevllm.operators.moe import MoeOpSpec, resolve_moe_provider -from sparsevllm.operators.moe_router import ( - MoeRouterOpSpec, - resolve_moe_router_provider, -) - - -WEIGHT_PREFIX = "model.language_model.layers.0.mlp." -WEIGHT_NAMES = ( - "gate.weight", - "experts.gate_up_proj", - "experts.down_proj", - "shared_expert.gate_proj.weight", - "shared_expert.up_proj.weight", - "shared_expert.down_proj.weight", - "shared_expert_gate.weight", -) - - -def _write_json(path: Path, value: Any) -> None: - path.write_text( - json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - for chunk in iter(lambda: handle.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _git_value(*args: str) -> str | None: - result = subprocess.run( - ["git", *args], - cwd=REPO_ROOT, - check=False, - capture_output=True, - text=True, - ) - value = result.stdout.strip() - return value or None - - -def _load_layer_weights(model_path: Path) -> dict[str, torch.Tensor]: - index_path = model_path / "model.safetensors.index.json" - if not index_path.is_file(): - raise FileNotFoundError(f"Missing checkpoint index: {index_path}") - weight_map = json.loads(index_path.read_text(encoding="utf-8"))["weight_map"] - loaded: dict[str, torch.Tensor] = {} - by_shard: dict[str, list[tuple[str, str]]] = {} - for local_name in WEIGHT_NAMES: - checkpoint_name = WEIGHT_PREFIX + local_name - shard = weight_map.get(checkpoint_name) - if shard is None: - raise KeyError(f"Missing checkpoint tensor {checkpoint_name!r}.") - by_shard.setdefault(shard, []).append((checkpoint_name, local_name)) - for shard, names in by_shard.items(): - shard_path = model_path / shard - if not shard_path.is_file(): - raise FileNotFoundError(f"Missing checkpoint shard: {shard_path}") - with safe_open(shard_path, framework="pt", device="cpu") as handle: - for checkpoint_name, local_name in names: - tensor = handle.get_tensor(checkpoint_name) - if tensor.dtype != torch.bfloat16: - raise TypeError( - f"{checkpoint_name} must be BF16, got {tensor.dtype}." - ) - loaded[local_name] = tensor - return loaded - - -def _tensor_summary(value: torch.Tensor) -> dict[str, Any]: - finite = torch.isfinite(value) - float_value = value.float() - return { - "shape": list(value.shape), - "dtype": str(value.dtype), - "finite": bool(finite.all().item()), - "min": float(float_value.min().item()), - "max": float(float_value.max().item()), - "mean": float(float_value.mean().item()), - } - - -def _error_metrics( - actual: torch.Tensor, - expected: torch.Tensor, -) -> dict[str, float]: - difference = (actual.float() - expected.float()).abs() - denominator = expected.float().abs().clamp_min(1.0e-6) - return { - "max_abs": float(difference.max().item()), - "mean_abs": float(difference.mean().item()), - "max_rel": float((difference / denominator).max().item()), - } - - -def _record_close( - name: str, - actual: torch.Tensor, - expected: torch.Tensor, - *, - atol: float, - rtol: float, -) -> dict[str, Any]: - metrics = _error_metrics(actual, expected) - passed = bool(torch.allclose(actual.float(), expected.float(), atol=atol, rtol=rtol)) - return { - "component": name, - "status": "success" if passed else "metric_failed", - "atol": float(atol), - "rtol": float(rtol), - **metrics, - } - - -def main() -> int: - parser = argparse.ArgumentParser( - description="Compare Qwen3.6 MoE BF16 checkpoint math with Transformers." - ) - parser.add_argument("--model", type=Path, required=True) - parser.add_argument("--output-dir", type=Path, required=True) - parser.add_argument("--tokens", type=int, default=17) - parser.add_argument("--seed", type=int, default=20260810) - args = parser.parse_args() - if args.tokens <= 0: - raise ValueError(f"--tokens must be positive, got {args.tokens}.") - if not torch.cuda.is_available(): - raise RuntimeError("This validation requires one CUDA device.") - - model_path = args.model.resolve() - output_dir = args.output_dir.resolve() - output_dir.mkdir(parents=True, exist_ok=False) - index_path = model_path / "model.safetensors.index.json" - config_path = model_path / "config.json" - run_info = { - "created_at": datetime.now().isoformat(timespec="seconds"), - "command": " ".join(sys.argv), - "model": str(model_path), - "model_config_sha256": _sha256(config_path), - "model_index_sha256": _sha256(index_path), - "seed": int(args.seed), - "tokens": int(args.tokens), - "dtype": "torch.bfloat16", - "device": torch.cuda.get_device_name(0), - "git_commit": _git_value("rev-parse", "HEAD"), - "git_branch": _git_value("branch", "--show-current"), - "git_dirty": bool(_git_value("status", "--porcelain")), - "tolerances": { - "routing_weights": {"atol": 0.004, "rtol": 0.004}, - "routed_experts": {"atol": 0.05, "rtol": 0.05}, - "shared_expert": {"atol": 0.02, "rtol": 0.02}, - "moe_output": {"atol": 0.05, "rtol": 0.05}, - }, - "requested_moe_provider": os.getenv("SPARSEVLLM_MOE_PROVIDER", "auto"), - "requested_moe_router_provider": os.getenv( - "SPARSEVLLM_MOE_ROUTER_PROVIDER", "auto" - ), - } - - weights = _load_layer_weights(model_path) - outer_config = AutoConfig.from_pretrained(model_path, local_files_only=True) - config = outer_config.text_config - if config.torch_dtype != torch.bfloat16: - raise TypeError(f"Reference config must be BF16, got {config.torch_dtype}.") - num_experts = int(config.num_experts) - expert_intermediate_size = int(weights["experts.gate_up_proj"].shape[1] // 2) - moe_spec = MoeOpSpec( - num_experts=num_experts, - num_local_experts=num_experts, - hidden_size=int(config.hidden_size), - intermediate_size=expert_intermediate_size, - top_k=int(config.num_experts_per_tok), - activation_dtype=torch.bfloat16, - weight_dtype=torch.bfloat16, - block_shape=None, - ep_size=1, - cuda_graph=False, - tp_size=1, - routing_method="softmax", - ) - moe_provider = resolve_moe_provider(moe_spec, device_index=0) - router_spec = MoeRouterOpSpec( - num_experts=num_experts, - top_k=int(config.num_experts_per_tok), - activation_dtype=torch.bfloat16, - norm_topk_prob=True, - cuda_graph=False, - ) - router_provider = resolve_moe_router_provider(router_spec, device_index=0) - run_info["resolved_moe_provider"] = moe_provider.name - run_info["resolved_moe_router_provider"] = router_provider.name - _write_json(output_dir / "run_info.json", run_info) - - previous_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.bfloat16) - try: - reference = Qwen3_5MoeSparseMoeBlock(config) - finally: - torch.set_default_dtype(previous_dtype) - missing, unexpected = reference.load_state_dict(weights, strict=False) - if missing or unexpected: - raise RuntimeError( - f"Transformers MoE weight mismatch: missing={missing}, unexpected={unexpected}." - ) - reference = reference.to(device="cuda", dtype=torch.bfloat16).eval() - weights = { - name: tensor.to(device="cuda", non_blocking=False) - for name, tensor in weights.items() - } - - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - hidden = torch.randn( - (args.tokens, int(config.hidden_size)), - device="cuda", - dtype=torch.bfloat16, - ) - with torch.inference_mode(): - reference_logits, reference_routing_weights, reference_ids = reference.gate( - hidden - ) - reference_routed = reference.experts( - hidden, reference_ids, reference_routing_weights - ) - reference_shared = reference.shared_expert(hidden) - reference_shared = torch.sigmoid(reference.shared_expert_gate(hidden)) * reference_shared - reference_output = reference_routed + reference_shared - - actual_logits = F.linear(hidden, weights["gate.weight"]) - actual_routing_weights, actual_ids = router_provider.run( - router_spec, - actual_logits, - ) - actual_routed = moe_provider.run( - moe_spec, - hidden, - actual_ids, - actual_routing_weights, - weights["experts.gate_up_proj"], - weights["experts.down_proj"], - None, - None, - local_expert_start=0, - ep_rank=0, - ) - shared_gate = F.linear(hidden, weights["shared_expert.gate_proj.weight"]) - shared_up = F.linear(hidden, weights["shared_expert.up_proj.weight"]) - actual_shared = F.linear( - F.silu(shared_gate) * shared_up, - weights["shared_expert.down_proj.weight"], - ) - actual_shared *= torch.sigmoid( - F.linear(hidden, weights["shared_expert_gate.weight"]) - ) - actual_output = actual_routed + actual_shared - torch.cuda.synchronize() - - records = [ - _record_close( - "router_logits", - actual_logits, - reference_logits, - atol=0.0, - rtol=0.0, - ), - { - "component": "routing_ids", - "status": "success" if torch.equal(actual_ids, reference_ids.to(torch.int32)) else "metric_failed", - "mismatch_count": int((actual_ids != reference_ids.to(torch.int32)).sum().item()), - }, - _record_close( - "routing_weights", - actual_routing_weights, - reference_routing_weights, - atol=0.004, - rtol=0.004, - ), - _record_close( - "routed_experts", - actual_routed, - reference_routed, - atol=0.05, - rtol=0.05, - ), - _record_close( - "shared_expert", - actual_shared, - reference_shared, - atol=0.02, - rtol=0.02, - ), - _record_close( - "moe_output", - actual_output, - reference_output, - atol=0.05, - rtol=0.05, - ), - ] - raw_outputs = { - "hidden": hidden.cpu(), - "reference_logits": reference_logits.cpu(), - "actual_logits": actual_logits.cpu(), - "reference_ids": reference_ids.cpu(), - "actual_ids": actual_ids.cpu(), - "reference_routing_weights": reference_routing_weights.cpu(), - "actual_routing_weights": actual_routing_weights.cpu(), - "reference_routed": reference_routed.cpu(), - "actual_routed": actual_routed.cpu(), - "reference_shared": reference_shared.cpu(), - "actual_shared": actual_shared.cpu(), - "reference_output": reference_output.cpu(), - "actual_output": actual_output.cpu(), - } - torch.save(raw_outputs, output_dir / "raw_outputs.pt") - parsed = { - name: _tensor_summary(tensor) - for name, tensor in raw_outputs.items() - } - _write_json(output_dir / "parsed_outputs.json", parsed) - with (output_dir / "per_sample_results.jsonl").open("w", encoding="utf-8") as handle: - for record in records: - handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n") - failed = [record for record in records if record["status"] != "success"] - aggregate = { - "status": "success" if not failed else "metric_failed", - "num_checks": len(records), - "success_checks": len(records) - len(failed), - "failed_checks": len(failed), - "records": records, - } - _write_json(output_dir / "aggregate_metrics.json", aggregate) - print(json.dumps(aggregate, ensure_ascii=False, indent=2)) - return 0 if not failed else 2 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/validation/qwen36_summarize_microbench.py b/scripts/validation/qwen36_summarize_microbench.py deleted file mode 100644 index 18b16b2c..00000000 --- a/scripts/validation/qwen36_summarize_microbench.py +++ /dev/null @@ -1,260 +0,0 @@ -from __future__ import annotations - -import argparse -import json -import statistics -import subprocess -import sys -from datetime import datetime -from pathlib import Path -from typing import Any - - -CASES = ( - (1024, 1, 512), - (4096, 8, 512), - (32768, 2, 512), - (65536, 1, 1024), - (131072, 1, 1024), -) -TOPOLOGIES = ("single", "pure_tp", "tp_ep") -MODES = ("graph", "eager") -METRICS = ( - "ttft_s", - "prefill_tok_s", - "decode_tok_s", - "itl_ms", - "end_to_end_tok_s", - "peak_memory_gb", -) - - -def _read_jsonl(path: Path) -> list[dict[str, Any]]: - rows = [] - for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): - value = json.loads(line) - if not isinstance(value, dict): - raise TypeError(f"Expected an object at {path}:{line_no}.") - rows.append(value) - return rows - - -def _write_json(path: Path, value: Any) -> None: - path.write_text( - json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - - -def _stats(rows: list[dict[str, Any]], metric: str) -> dict[str, float]: - values = [float(row[metric]) for row in rows] - return { - "median": statistics.median(values), - "min": min(values), - "max": max(values), - "relative_range": ( - (max(values) - min(values)) / statistics.median(values) - if statistics.median(values) - else 0.0 - ), - } - - -def main() -> int: - parser = argparse.ArgumentParser( - description="Validate and summarize a completed Qwen3.6 microbench matrix." - ) - parser.add_argument("--matrix-dir", type=Path, required=True) - parser.add_argument("--output-dir", type=Path, required=True) - parser.add_argument("--repeats", type=int, default=2) - parser.add_argument( - "--modes", - default="graph,eager", - help="Comma-separated subset of graph,eager.", - ) - args = parser.parse_args() - modes = tuple(part.strip() for part in args.modes.split(",") if part.strip()) - if not modes or len(set(modes)) != len(modes) or any( - mode not in MODES for mode in modes - ): - raise ValueError( - "--modes must contain each of 'graph' and 'eager' at most once, " - f"got {args.modes!r}." - ) - matrix_dir = args.matrix_dir.resolve() - output_dir = args.output_dir.resolve() - output_dir.mkdir(parents=True, exist_ok=False) - records = _read_jsonl(matrix_dir / "per_sample_results.jsonl") - - expected_keys = { - (topology, mode, length, batch, output) - for length, batch, output in CASES - for topology in TOPOLOGIES - for mode in modes - } - groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {} - validation_records = [] - for record in records: - key = ( - record.get("topology"), - record.get("mode"), - int(record.get("prompt_tokens")), - int(record.get("batch_size")), - int(record.get("max_new_tokens")), - ) - groups.setdefault(key, []).append(record) - workers = record.get("worker_runtime_status") or [] - expected_workers = 1 if record.get("topology") == "single" else 2 - graph_active = ( - record.get("mode") != "graph" - or ( - len(workers) == expected_workers - and all(worker.get("decode_cuda_graph_active") for worker in workers) - ) - ) - providers_ok = ( - record.get("moe_expert_provider") == "triton" - and record.get("moe_router_provider") == "triton" - and all( - worker.get("moe_expert_provider") == "triton" - and worker.get("moe_router_provider") == "triton" - for worker in workers - ) - ) - validation_records.append( - { - "check": "measured_run", - "run_name": record.get("run_name"), - "status": ( - "success" - if record.get("status") == "success" - and graph_active - and providers_ok - else "metric_failed" - ), - "graph_active_on_all_workers": graph_active, - "providers_ok": providers_ok, - "worker_count": len(workers), - } - ) - - missing_groups = sorted(expected_keys - set(groups)) - unexpected_groups = sorted(set(groups) - expected_keys) - summaries = [] - for key in sorted(expected_keys): - rows = groups.get(key, []) - successful = [row for row in rows if row.get("status") == "success"] - summaries.append( - { - "topology": key[0], - "mode": key[1], - "prompt_tokens": key[2], - "batch_size": key[3], - "max_new_tokens": key[4], - "status": ( - "success" - if len(successful) == args.repeats - else "metric_failed" - ), - "successful_repeats": len(successful), - "required_repeats": args.repeats, - "metrics": ( - {metric: _stats(successful, metric) for metric in METRICS} - if len(successful) == args.repeats - else None - ), - } - ) - - summary_by_key = { - ( - row["topology"], - row["mode"], - row["prompt_tokens"], - row["batch_size"], - row["max_new_tokens"], - ): row - for row in summaries - } - speedups = [] - if {"graph", "eager"} <= set(modes): - for length, batch, output in CASES: - for topology in TOPOLOGIES: - graph = summary_by_key[(topology, "graph", length, batch, output)] - eager = summary_by_key[(topology, "eager", length, batch, output)] - if graph["status"] != "success" or eager["status"] != "success": - speedups.append( - { - "topology": topology, - "prompt_tokens": length, - "batch_size": batch, - "max_new_tokens": output, - "status": "metric_failed", - } - ) - continue - graph_decode = graph["metrics"]["decode_tok_s"]["median"] - eager_decode = eager["metrics"]["decode_tok_s"]["median"] - graph_itl = graph["metrics"]["itl_ms"]["median"] - eager_itl = eager["metrics"]["itl_ms"]["median"] - speedups.append( - { - "topology": topology, - "prompt_tokens": length, - "batch_size": batch, - "max_new_tokens": output, - "status": "success", - "decode_throughput_speedup": graph_decode / eager_decode, - "itl_speedup": eager_itl / graph_itl, - "graph_decode_tok_s": graph_decode, - "eager_decode_tok_s": eager_decode, - } - ) - - failed = [ - record - for record in [*validation_records, *summaries, *speedups] - if record["status"] != "success" - ] - aggregate = { - "status": ( - "success" - if not failed and not missing_groups and not unexpected_groups - else "metric_failed" - ), - "num_runs": len(records), - "expected_runs": len(expected_keys) * args.repeats, - "missing_groups": missing_groups, - "unexpected_groups": unexpected_groups, - "failed_checks": len(failed), - "summaries": summaries, - "graph_speedups": speedups, - } - run_info = { - "created_at": datetime.now().isoformat(timespec="seconds"), - "command": " ".join(sys.argv), - "git_commit": subprocess.run( - ["git", "rev-parse", "HEAD"], text=True, capture_output=True, check=True - ).stdout.strip(), - "matrix_dir": str(matrix_dir), - "repeats": args.repeats, - "modes": list(modes), - } - _write_json(output_dir / "run_info.json", run_info) - _write_json(output_dir / "raw_outputs.json", {"matrix_dir": str(matrix_dir)}) - _write_json( - output_dir / "parsed_outputs.json", - {"run_checks": validation_records, "summaries": summaries, "speedups": speedups}, - ) - with (output_dir / "per_sample_results.jsonl").open( - "w", encoding="utf-8" - ) as handle: - for record in [*validation_records, *summaries, *speedups]: - handle.write(json.dumps(record, ensure_ascii=False) + "\n") - _write_json(output_dir / "aggregate_metrics.json", aggregate) - print(json.dumps(aggregate, ensure_ascii=False, indent=2)) - return 0 if aggregate["status"] == "success" else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/validation/run_qwen36_microbench_matrix.py b/scripts/validation/run_qwen36_microbench_matrix.py deleted file mode 100644 index 79ce3dd3..00000000 --- a/scripts/validation/run_qwen36_microbench_matrix.py +++ /dev/null @@ -1,391 +0,0 @@ -from __future__ import annotations - -import argparse -import hashlib -import json -import os -import statistics -import subprocess -import sys -import time -from datetime import datetime -from pathlib import Path -from typing import Any - - -CASES = ( - (1024, 1, 512), - (4096, 8, 512), - (32768, 2, 512), - (65536, 1, 1024), - (131072, 1, 1024), -) -TOPOLOGIES = ( - ("single", "0", 1, 1), - ("pure_tp", "0,1", 2, 1), - ("tp_ep", "0,1", 2, 2), -) -MODES = ("graph", "eager") -METRICS = ( - "ttft_s", - "prefill_tok_s", - "decode_tok_s", - "itl_ms", - "end_to_end_tok_s", - "peak_memory_gb", -) - - -def _write_json(path: Path, value: Any) -> None: - path.write_text( - json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - for chunk in iter(lambda: handle.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _gpu_snapshot() -> dict[str, Any]: - processes = subprocess.run( - [ - "nvidia-smi", - "--query-compute-apps=gpu_uuid,pid,process_name,used_memory", - "--format=csv,noheader,nounits", - ], - text=True, - capture_output=True, - check=True, - ).stdout.strip() - gpu_lines = subprocess.run( - [ - "nvidia-smi", - "--query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu", - "--format=csv,noheader,nounits", - ], - text=True, - capture_output=True, - check=True, - ).stdout.strip().splitlines() - gpus = [] - for line in gpu_lines: - fields = [part.strip() for part in line.split(",")] - if len(fields) != 6: - raise RuntimeError(f"Unexpected nvidia-smi GPU row: {line!r}.") - gpus.append( - { - "index": int(fields[0]), - "uuid": fields[1], - "name": fields[2], - "memory_used_mib": int(fields[3]), - "memory_total_mib": int(fields[4]), - "utilization_percent": int(fields[5]), - } - ) - return {"compute_processes": processes.splitlines() if processes else [], "gpus": gpus} - - -def _wait_for_all_devices_idle(timeout_s: int) -> dict[str, Any]: - deadline = time.monotonic() + timeout_s - while True: - snapshot = _gpu_snapshot() - idle = not snapshot["compute_processes"] and all( - gpu["utilization_percent"] <= 1 for gpu in snapshot["gpus"] - ) - if idle: - return snapshot - if time.monotonic() >= deadline: - raise TimeoutError( - f"GPUs remained busy for {timeout_s}s; last snapshot={snapshot!r}." - ) - print(f"[gpu-idle] busy; waiting 10s: {snapshot}", flush=True) - time.sleep(10) - - -def _read_single_record(run_dir: Path) -> dict[str, Any]: - aggregate_path = run_dir / "aggregate_metrics.json" - if not aggregate_path.is_file(): - return { - "status": "model_failed", - "error": f"missing artifact {aggregate_path}", - } - aggregate = json.loads(aggregate_path.read_text(encoding="utf-8")) - records = aggregate.get("records") - if not isinstance(records, list) or len(records) != 1: - return { - "status": "model_failed", - "error": f"expected exactly one microbench record, got {records!r}", - } - return records[0] - - -def _aggregate(records: list[dict[str, Any]], repeats: int) -> dict[str, Any]: - groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {} - for record in records: - key = ( - record["topology"], - record["mode"], - record["prompt_tokens"], - record["batch_size"], - record["max_new_tokens"], - ) - groups.setdefault(key, []).append(record) - summaries = [] - failed_groups = 0 - for key, rows in sorted(groups.items()): - successful = [row for row in rows if row.get("status") == "success"] - status = "success" if len(successful) == repeats else "model_failed" - if status != "success": - failed_groups += 1 - metrics = {} - for metric in METRICS: - values = [float(row[metric]) for row in successful if metric in row] - metrics[metric] = ( - None - if len(values) != repeats - else { - "median": statistics.median(values), - "min": min(values), - "max": max(values), - } - ) - summaries.append( - { - "topology": key[0], - "mode": key[1], - "prompt_tokens": key[2], - "batch_size": key[3], - "max_new_tokens": key[4], - "status": status, - "successful_repeats": len(successful), - "required_repeats": repeats, - "metrics": metrics, - } - ) - return { - "status": "success" if not failed_groups else "model_failed", - "num_groups": len(summaries), - "successful_groups": len(summaries) - failed_groups, - "failed_groups": failed_groups, - "required_repeats": repeats, - "summaries": summaries, - } - - -def main() -> int: - parser = argparse.ArgumentParser(description="Run the fixed Qwen3.6 MoE benchmark matrix.") - parser.add_argument("--model", type=Path, required=True) - parser.add_argument("--output-dir", type=Path, required=True) - parser.add_argument("--repeats", type=int, default=2) - parser.add_argument( - "--modes", - default="graph,eager", - help="Comma-separated subset of graph,eager.", - ) - parser.add_argument("--idle-timeout", type=int, default=600) - parser.add_argument("--run-timeout", type=int, default=3600) - args = parser.parse_args() - if args.repeats < 2: - raise ValueError("The acceptance matrix requires at least two measured repeats.") - if args.run_timeout <= 0: - raise ValueError("--run-timeout must be positive.") - modes = tuple(part.strip() for part in args.modes.split(",") if part.strip()) - if not modes or len(set(modes)) != len(modes) or any( - mode not in MODES for mode in modes - ): - raise ValueError( - "--modes must contain each of 'graph' and 'eager' at most once, " - f"got {args.modes!r}." - ) - model = args.model.resolve() - output_dir = args.output_dir.resolve() - output_dir.mkdir(parents=True, exist_ok=False) - for name in ("config.json", "model.safetensors.index.json"): - if not (model / name).is_file(): - raise FileNotFoundError(f"Required model metadata is missing: {model / name}.") - - initial_gpu = _wait_for_all_devices_idle(args.idle_timeout) - run_info = { - "created_at": datetime.now().isoformat(timespec="seconds"), - "command": " ".join([sys.executable, *sys.argv]), - "model": str(model), - "model_config_sha256": _sha256(model / "config.json"), - "model_index_sha256": _sha256(model / "model.safetensors.index.json"), - "git_commit": subprocess.run( - ["git", "rev-parse", "HEAD"], text=True, capture_output=True, check=True - ).stdout.strip(), - "git_branch": subprocess.run( - ["git", "branch", "--show-current"], text=True, capture_output=True, check=True - ).stdout.strip(), - "git_dirty": bool( - subprocess.run( - ["git", "status", "--porcelain"], - text=True, - capture_output=True, - check=True, - ).stdout.strip() - ), - "cases": [ - {"prompt_tokens": length, "batch_size": batch, "max_new_tokens": output} - for length, batch, output in CASES - ], - "topologies": [ - { - "name": name, - "cuda_visible_devices": devices, - "outer_tp_size": tp, - "expert_parallel_size": ep, - "moe_tp_size": tp // ep, - } - for name, devices, tp, ep in TOPOLOGIES - ], - "modes": list(modes), - "repeats": args.repeats, - "seed": 20260810, - "warmup_output_len": 8, - "run_timeout_s": args.run_timeout, - "initial_gpu_snapshot": initial_gpu, - } - _write_json(output_dir / "run_info.json", run_info) - - records: list[dict[str, Any]] = [] - total = len(CASES) * len(TOPOLOGIES) * len(modes) * args.repeats - run_index = 0 - for length, batch, max_new_tokens in CASES: - for topology, devices, tp_size, ep_size in TOPOLOGIES: - for mode in modes: - for repeat in range(1, args.repeats + 1): - run_index += 1 - idle_snapshot = _wait_for_all_devices_idle(args.idle_timeout) - run_name = ( - f"{length}_bs{batch}_out{max_new_tokens}_" - f"{topology}_{mode}_r{repeat}" - ) - run_dir = output_dir / "runs" / run_name - hyper_params = { - "tensor_parallel_size": tp_size, - "expert_parallel_size": ep_size, - "enforce_eager": mode == "eager", - "decode_cuda_graph": mode == "graph", - "gpu_memory_utilization": 0.9, - "engine_prefill_chunk_size": 8192, - "max_num_batched_tokens": 65536, - "weight_loading_workers": 16, - } - command = [ - sys.executable, - "benchmark/microbench.py", - "--model_path", - str(model), - "--lengths", - str(length), - "--batch_sizes", - str(batch), - "--output_len", - str(max_new_tokens), - "--methods", - "vanilla", - "--temperature", - "0", - "--top_p", - "1", - "--seed", - "20260810", - "--synchronize_step_timing", - "--warmup_output_len", - "8", - "--hyper_params", - json.dumps(hyper_params, separators=(",", ":")), - "--output_dir", - str(run_dir), - ] - env = os.environ.copy() - env.update( - { - "CUDA_VISIBLE_DEVICES": devices, - "PYTHONUNBUFFERED": "1", - "SPARSEVLLM_MOE_PROVIDER": "triton", - "SPARSEVLLM_MOE_ROUTER_PROVIDER": "triton", - } - ) - print( - f"[matrix {run_index}/{total}] {run_name}: {' '.join(command)}", - flush=True, - ) - timed_out = False - try: - process = subprocess.run( - command, - cwd=Path(__file__).resolve().parents[2], - env=env, - text=True, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - check=False, - timeout=args.run_timeout, - ) - console_output = process.stdout - returncode = process.returncode - except subprocess.TimeoutExpired as exc: - timed_out = True - console_output = exc.stdout or "" - if isinstance(console_output, bytes): - console_output = console_output.decode( - "utf-8", errors="replace" - ) - console_output += ( - "\nMatrix runner timed out this run after " - f"{args.run_timeout}s.\n" - ) - returncode = None - run_dir.mkdir(parents=True, exist_ok=True) - (run_dir / "console.log").write_text( - console_output, encoding="utf-8" - ) - record = _read_single_record(run_dir) - record.update( - { - "run_name": run_name, - "topology": topology, - "mode": mode, - "repeat": repeat, - "prompt_tokens": length, - "batch_size": batch, - "max_new_tokens": max_new_tokens, - "returncode": returncode, - "timed_out": timed_out, - "command": command, - "idle_snapshot_before": idle_snapshot, - "artifact_dir": str(run_dir), - } - ) - if timed_out or returncode != 0: - record["status"] = "model_failed" - records.append(record) - with (output_dir / "per_sample_results.jsonl").open( - "a", encoding="utf-8" - ) as handle: - handle.write(json.dumps(record, ensure_ascii=False) + "\n") - print( - f"[matrix {run_index}/{total}] status={record['status']}", - flush=True, - ) - - aggregate = _aggregate(records, args.repeats) - _write_json(output_dir / "aggregate_metrics.json", aggregate) - _write_json(output_dir / "parsed_outputs.json", {"records": records}) - _write_json( - output_dir / "raw_outputs.json", - {"run_artifact_dirs": [record["artifact_dir"] for record in records]}, - ) - print(json.dumps(aggregate, ensure_ascii=False, indent=2), flush=True) - return 0 if aggregate["status"] == "success" else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/src/sparsevllm/engine/llm_engine.py b/src/sparsevllm/engine/llm_engine.py index 53394325..86cf386a 100644 --- a/src/sparsevllm/engine/llm_engine.py +++ b/src/sparsevllm/engine/llm_engine.py @@ -31,7 +31,10 @@ RequestAdmission, stable_token_digest, ) -from sparsevllm.method_registry import normalize_sparse_method +from sparsevllm.method_registry import ( + OUTER_TP_MOE_MODEL_TYPES, + normalize_sparse_method, +) from sparsevllm.utils.profiler import profiler def _deltakv_graph_warmup_profile(config: Config) -> str: @@ -62,7 +65,9 @@ def _use_graph_scaled_warmup(config: Config) -> bool: def _moe_workspace_warmup_token_counts(config: Config) -> tuple[int, ...]: - if int(getattr(config.hf_config, "num_experts", 0) or 0) <= 0: + model_type = str(getattr(config.hf_config, "model_type", "") or "") + has_experts = int(getattr(config.hf_config, "num_experts", 0) or 0) > 0 + if not has_experts and model_type not in OUTER_TP_MOE_MODEL_TYPES: return () max_batched_tokens = int(config.max_num_batched_tokens) diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 1b84f32d..4fe06145 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -88,7 +88,6 @@ "refresh_prefix_cache_hit", "reset_after_warmup", "run", - "runtime_diagnostic_status", "set_warmup_fake_prefill_attention", "warmup_moe_workspace", } @@ -119,9 +118,6 @@ def __init__( profiler.set_enabled(config.enable_profiler and rank == 0) hf_config = config.hf_config self.enforce_eager = config.enforce_eager - self.debug_runtime_enabled = ( - os.getenv("SPARSEVLLM_DEBUG_RUNTIME", "0") == "1" - ) self.world_size = config.world_size self.rank = rank self.event = event @@ -846,14 +842,12 @@ def debug_sparse_state_summary(self) -> dict[str, object]: } def debug_last_logits_cpu(self) -> torch.Tensor | None: - if self.rank != 0: - return None logits = getattr(self, "debug_last_logits", None) if logits is None: raise RuntimeError( "No debug logits are available. Set SPARSEVLLM_DEBUG_RUNTIME=1 before engine startup." ) - return logits.detach().cpu() + return logits.detach().cpu() if self.rank == 0 else None def debug_hidden_states_cpu(self) -> dict[int, torch.Tensor] | None: model = getattr(getattr(self, "model", None), "model", None) @@ -937,15 +931,13 @@ def _debug_any_mismatch_from_world_rank_zero(self, tensor: torch.Tensor) -> bool def debug_replica_consistency(self) -> dict[str, object] | None: logits = getattr(self, "debug_last_logits", None) - if self.world_size == 1 and logits is None: + if logits is None: return None - if self.world_size == 1: - logits_max_abs, logits_tolerance_ratio = 0.0, 0.0 - else: - # Only world rank 0 materializes LM-head logits under tensor - # parallelism. Cross-rank consistency is checked on the synchronized - # MoE outputs below instead. - logits_max_abs, logits_tolerance_ratio = None, None + logits_max_abs, logits_tolerance_ratio = self._debug_float_error_from_world_rank_zero( + logits, + atol=0.05, + rtol=0.05, + ) result: dict[str, object] = { "last_logits_max_abs": logits_max_abs, "last_logits_tolerance_ratio": logits_tolerance_ratio, @@ -1003,62 +995,6 @@ def debug_sparse_state_summaries(self) -> list[dict[str, object]] | None: ) return summaries if self.rank == 0 else None - def runtime_diagnostic_status(self) -> list[dict[str, object]] | None: - graph_states = getattr(self.decode_cuda_graph_runner, "_graphs", {}) - graph_count = sum( - 1 - for state in graph_states.values() - if getattr(state, "graph", None) is not None - ) - local_status = { - "world_rank": int(self.parallel_context.world_rank), - "attention_tp_rank": int(self.parallel_context.attention_tp_rank), - "attention_tp_size": int(self.parallel_context.attention_tp_size), - "moe_tp_rank": int(self.parallel_context.moe_tp_rank), - "moe_tp_size": int(self.parallel_context.moe_tp_size), - "ep_rank": int(self.parallel_context.ep_rank), - "ep_size": int(self.parallel_context.ep_size), - "decode_cuda_graph_configured": bool(self.config.decode_cuda_graph), - "decode_cuda_graph_state_count": int(len(graph_states)), - "decode_cuda_graph_graph_count": int(graph_count), - "decode_cuda_graph_active": bool( - self.config.decode_cuda_graph and graph_count > 0 - ), - } - model_status = getattr(self.model, "runtime_diagnostic_status", None) - if callable(model_status): - try: - extra_status = model_status() - if not isinstance(extra_status, dict): - raise TypeError( - "Model runtime_diagnostic_status() must return a dict, " - f"got {extra_status!r}." - ) - local_status.update(extra_status) - except Exception as exc: - if self.world_size == 1: - raise - local_status["runtime_diagnostic_error"] = repr(exc) - if self.world_size == 1: - return [local_status] - statuses = [None] * self.world_size - dist.all_gather_object( - statuses, - local_status, - group=self.parallel_context.world.process_group, - ) - errors = [ - status.get("runtime_diagnostic_error") - for status in statuses - if status.get("runtime_diagnostic_error") is not None - ] - if errors: - raise RuntimeError( - "Model runtime diagnostics failed on at least one worker: " - f"{errors}." - ) - return statuses if self.rank == 0 else None - def _long_text_threshold(self, is_prefill: bool) -> int: del is_prefill if self.config.vllm_sparse_method in ("streamingllm", "attention-sink", "attention_sink"): @@ -1253,20 +1189,14 @@ def set_decode_cuda_graph_max_context_len_override(self, max_context_len: int | def set_omnikv_decode_graph_max_context_len_override(self, max_context_len: int | None): self.set_decode_cuda_graph_max_context_len_override(max_context_len) - def _capture_debug_logits(self, logits: torch.Tensor | None) -> None: - if ( - getattr(self, "debug_runtime_enabled", False) - and isinstance(logits, torch.Tensor) - ): - self.debug_last_logits = logits.detach().clone() - @torch.inference_mode() def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool): """物理执行逻辑：统一使用 Eager 模式""" _stage = 'prefill' if is_prefill else 'decode' with profiler.record(f"model_run_model_{_stage}"): logits = self.model.compute_logits(self.model(input_ids, positions)) - self._capture_debug_logits(logits) + if os.getenv("SPARSEVLLM_DEBUG_RUNTIME", "0") == "1": + self.debug_last_logits = logits.detach().clone() return logits def run_logits_for_compare(self, seqs: list[Sequence], is_prefill: bool) -> torch.Tensor | None: @@ -1347,7 +1277,6 @@ def run( else: logits = self.decode_cuda_graph_runner.run_eager_static(seqs) graph_token_ids = None - self._capture_debug_logits(logits) if self.rank != 0: self._post_sparse_forward(seqs, is_prefill) return None, None diff --git a/src/sparsevllm/models/qwen3_5_moe.py b/src/sparsevllm/models/qwen3_5_moe.py index df995fe6..db2a8cb6 100644 --- a/src/sparsevllm/models/qwen3_5_moe.py +++ b/src/sparsevllm/models/qwen3_5_moe.py @@ -1,6 +1,5 @@ from __future__ import annotations -import os import re import torch @@ -71,12 +70,12 @@ def __init__(self, config) -> None: def forward( self, hidden_states: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: router_logits = F.linear(hidden_states, self.weight) topk_weights, topk_ids = self.provider.run( self.op_spec, router_logits ) - return router_logits, topk_weights, topk_ids + return topk_weights, topk_ids class Qwen35MoePackedExperts(Qwen3MoePackedExperts): @@ -224,7 +223,6 @@ class Qwen35MoeSparseMoeBlock(nn.Module): def __init__(self, config) -> None: super().__init__() self.parallel_context = get_parallel_context() - self.debug_enabled = os.getenv("SPARSEVLLM_DEBUG_MOE", "0") == "1" self.mlp_chunk_size = int(getattr(config, "mlp_chunk_size", 16384)) if self.mlp_chunk_size <= 0: raise ValueError( @@ -243,28 +241,14 @@ def __init__(self, config) -> None: def _forward_chunk( self, hidden_states: torch.Tensor, - ) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - ]: + ) -> torch.Tensor: shared_output = self.shared_expert(hidden_states) - router_logits, topk_weights, topk_ids = self.gate(hidden_states) + topk_weights, topk_ids = self.gate(hidden_states) local_output = self.experts(hidden_states, topk_ids, topk_weights) routed_output = self.parallel_context.world_all_reduce(local_output) shared_gate = torch.sigmoid(self.shared_expert_gate(hidden_states)) gated_shared_output = shared_gate * shared_output - return ( - routed_output + gated_shared_output, - router_logits, - topk_weights, - topk_ids, - local_output, - gated_shared_output, - ) + return routed_output + gated_shared_output def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if hidden_states.dim() != 2: @@ -272,62 +256,11 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: "Qwen35MoeSparseMoeBlock expects [tokens, hidden], " f"got {tuple(hidden_states.shape)}." ) - debug_enabled = self.debug_enabled - if debug_enabled: - self.debug_last_input = hidden_states.detach().clone() - chunks = hidden_states.split(self.mlp_chunk_size, dim=0) outputs = [] - router_logits_chunks = [] - topk_weights_chunks = [] - topk_ids_chunks = [] - local_output_chunks = [] - shared_output_chunks = [] for chunk in chunks: - ( - output, - router_logits, - topk_weights, - topk_ids, - local_output, - shared_output, - ) = self._forward_chunk(chunk) - outputs.append(output) - if debug_enabled: - router_logits_chunks.append(router_logits) - topk_weights_chunks.append(topk_weights) - topk_ids_chunks.append(topk_ids) - local_output_chunks.append(local_output) - shared_output_chunks.append(shared_output) - output = outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0) - - if debug_enabled: - self.debug_last_router_logits = torch.cat( - router_logits_chunks, dim=0 - ).detach().clone() - self.debug_last_topk_weights = torch.cat( - topk_weights_chunks, dim=0 - ).detach().clone() - self.debug_last_topk_ids = torch.cat( - topk_ids_chunks, dim=0 - ).detach().clone() - self.debug_last_local_output = torch.cat( - local_output_chunks, dim=0 - ).detach().clone() - self.debug_last_shared_output = torch.cat( - shared_output_chunks, dim=0 - ).detach().clone() - local_mask = ( - self.debug_last_topk_ids >= self.experts.local_expert_start - ) & (self.debug_last_topk_ids < self.experts.local_expert_end) - local_hit_count = local_mask.sum() - self.debug_last_local_hit_count = ( - local_hit_count - if device_runtime.is_stream_capturing() - else int(local_hit_count.item()) - ) - self.debug_last_output = output.detach().clone() - return output + outputs.append(self._forward_chunk(chunk)) + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0) class Qwen35MoeDecoderLayer(Qwen35DecoderLayer): @@ -388,23 +321,6 @@ def recurrent_state_spec(config, attention_tp_size: int) -> RecurrentStateSpec: ), ) - def runtime_diagnostic_status(self) -> dict[str, object]: - experts = self.model.layers[0].mlp.experts - router = self.model.layers[0].mlp.gate - shared_linear = self.model.layers[0].mlp.shared_expert.gate_up_proj - return { - "moe_expert_provider": experts.provider.name, - "moe_router_provider": router.provider.name, - "moe_weight_dtype": str(experts.w13_weight.dtype), - "fp8_linear_provider": ( - shared_linear.quant_provider.name - if shared_linear.quantized - else None - ), - "local_expert_start": int(experts.local_expert_start), - "local_expert_end": int(experts.local_expert_end), - } - @torch.inference_mode() def warmup_moe(self, num_tokens: int = 1) -> None: num_tokens = int(num_tokens) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index d9530acd..e0d212fd 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -1,11 +1,9 @@ from __future__ import annotations -import os from dataclasses import dataclass from importlib.util import find_spec import torch -import torch.nn.functional as F import sparsevllm.platforms as platforms from sparsevllm.operators.registry import ( @@ -178,69 +176,6 @@ def run( MOE_REGISTRY: OpRegistry[MoeOpSpec, MoeProvider] = OpRegistry("routed MoE") -@MOE_REGISTRY.register -class TorchMoeProvider(MoeProvider): - """Clear eager-only CUDA reference for routed-expert semantics.""" - - name = "torch" - priority = 0 - gate_up_order = "gate_up" - - @classmethod - def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: - if caps.platform != PlatformEnum.CUDA: - return SupportResult.no(f"requires CUDA, got {caps.platform.name}") - if spec.cuda_graph: - return SupportResult.no("reference provider is eager-only") - if spec.weight_dtype != spec.activation_dtype: - return SupportResult.no( - "reference provider requires unquantized weights matching activations" - ) - if spec.block_shape is not None: - return SupportResult.no("reference provider does not support quantized weights") - 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, - ): - del ep_rank - if w13_scale_inv is not None or w2_scale_inv is not None: - raise RuntimeError("Torch MoE reference does not accept expert scales.") - output = torch.zeros_like(hidden_states) - for local_expert_id in range(spec.num_local_experts): - global_expert_id = int(local_expert_start) + local_expert_id - token_indices, topk_slots = torch.where( - topk_ids == global_expert_id - ) - if token_indices.numel() == 0: - continue - projected = F.linear( - hidden_states[token_indices], - w13_weight[local_expert_id], - ) - gate, up = projected.chunk(2, dim=-1) - expert_output = F.linear( - F.silu(gate) * up, - w2_weight[local_expert_id], - ) - routed = expert_output * topk_weights[ - token_indices, topk_slots - ].unsqueeze(-1) - output.index_add_(0, token_indices, routed) - return output - - @MOE_REGISTRY.register class TritonMinimaxM2FusedMoeProvider(MoeProvider): name = "triton_minimax_m2_fused" @@ -578,22 +513,4 @@ def resolve_moe_provider( 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)) - requested = os.getenv("SPARSEVLLM_MOE_PROVIDER", "auto").strip().lower() - if requested == "auto": - return OpResolver(MOE_REGISTRY).resolve(spec, caps).provider - - providers = {provider.name: provider for provider in MOE_REGISTRY.providers} - if requested not in providers: - choices = ", ".join(["auto", *sorted(providers)]) - raise ValueError( - "SPARSEVLLM_MOE_PROVIDER must be one of " - f"{choices}, got {requested!r}." - ) - provider_cls = providers[requested] - support = provider_cls.supports(spec, caps) - if not support.supported: - raise RuntimeError( - f"Requested MoE provider {requested!r} does not support " - f"spec={spec!r} on device={caps.device_name!r}: {support.reason}." - ) - return provider_cls() + return OpResolver(MOE_REGISTRY).resolve(spec, caps).provider diff --git a/src/sparsevllm/operators/moe_router.py b/src/sparsevllm/operators/moe_router.py index 57334719..fc8e74c3 100644 --- a/src/sparsevllm/operators/moe_router.py +++ b/src/sparsevllm/operators/moe_router.py @@ -1,6 +1,5 @@ from __future__ import annotations -import os from dataclasses import dataclass import torch @@ -95,35 +94,6 @@ def run( ) -@MOE_ROUTER_REGISTRY.register -class TorchMoeRouterProvider(MoeRouterProvider): - name = "torch" - priority = 0 - - @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 spec.cuda_graph: - return SupportResult.no("reference provider is eager-only") - return SupportResult.yes() - - def run( - self, - spec: MoeRouterOpSpec, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - probabilities = torch.softmax(router_logits, dim=-1, dtype=torch.float32) - weights, ids = torch.topk(probabilities, spec.top_k, dim=-1) - if spec.norm_topk_prob: - weights = weights / weights.sum(dim=-1, keepdim=True) - return weights.to(dtype=router_logits.dtype), ids.to(dtype=torch.int32) - - def resolve_moe_router_provider( spec: MoeRouterOpSpec, *, @@ -133,24 +103,4 @@ def resolve_moe_router_provider( 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)) - requested = os.getenv("SPARSEVLLM_MOE_ROUTER_PROVIDER", "auto").strip().lower() - if requested == "auto": - return OpResolver(MOE_ROUTER_REGISTRY).resolve(spec, caps).provider - - providers = { - provider.name: provider for provider in MOE_ROUTER_REGISTRY.providers - } - if requested not in providers: - choices = ", ".join(["auto", *sorted(providers)]) - raise ValueError( - "SPARSEVLLM_MOE_ROUTER_PROVIDER must be one of " - f"{choices}, got {requested!r}." - ) - provider_cls = providers[requested] - support = provider_cls.supports(spec, caps) - if not support.supported: - raise RuntimeError( - f"Requested MoE router provider {requested!r} does not support " - f"spec={spec!r} on device={caps.device_name!r}: {support.reason}." - ) - return provider_cls() + return OpResolver(MOE_ROUTER_REGISTRY).resolve(spec, caps).provider diff --git a/src/sparsevllm/triton_kernel/moe_topk.py b/src/sparsevllm/triton_kernel/moe_topk.py index 0bdc2a4a..ecfec960 100644 --- a/src/sparsevllm/triton_kernel/moe_topk.py +++ b/src/sparsevllm/triton_kernel/moe_topk.py @@ -205,8 +205,8 @@ def topk_softmax( 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 in {128, 256} " - "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}." ) diff --git a/tests/test_input_processor.py b/tests/test_input_processor.py deleted file mode 100644 index d78f4dc3..00000000 --- a/tests/test_input_processor.py +++ /dev/null @@ -1,54 +0,0 @@ -from unittest.mock import Mock - -import pytest - -from sparsevllm.engine.input_processor import tokenize_text_prompt - - -def test_tokenize_text_prompt_preserves_token_ids(): - tokenizer = Mock() - - assert tokenize_text_prompt(tokenizer, [1, 2, 3]) == [1, 2, 3] - tokenizer.encode.assert_not_called() - - -@pytest.mark.parametrize( - "prompt", - [ - {"image": "example.png", "text": "describe"}, - {"video": "example.mp4", "text": "describe"}, - {"mtp": True, "text": "continue"}, - [1, "2"], - ], -) -def test_tokenize_text_prompt_rejects_structured_inputs(prompt): - with pytest.raises(TypeError, match="text prompts only"): - tokenize_text_prompt(Mock(), prompt) - - -def test_tokenize_text_prompt_rejects_empty_token_ids(): - with pytest.raises(ValueError, match="must not be empty"): - tokenize_text_prompt(Mock(), []) - - -@pytest.mark.parametrize( - ("bos_token", "prompt", "add_special_tokens"), - [ - ("", "hello", True), - ("", "hello", False), - (None, "hello", False), - ], -) -def test_tokenize_text_prompt_handles_bos_once( - bos_token, - prompt, - add_special_tokens, -): - tokenizer = Mock(bos_token=bos_token) - tokenizer.encode.return_value = [1] - - assert tokenize_text_prompt(tokenizer, prompt) == [1] - tokenizer.encode.assert_called_once_with( - prompt, - add_special_tokens=add_special_tokens, - ) diff --git a/tests/test_longbench_deltakv_contracts.py b/tests/test_longbench_deltakv_contracts.py index eaba245e..1941209d 100644 --- a/tests/test_longbench_deltakv_contracts.py +++ b/tests/test_longbench_deltakv_contracts.py @@ -60,48 +60,6 @@ def test_longbench_data_validation_requires_explicit_root(self): finally: longbench_pred.DATA_PREFIX_PATH = old_root - def test_longbench_writes_prompt_and_canonical_per_sample_artifact(self): - with tempfile.TemporaryDirectory() as tmp: - task_path = str(Path(tmp) / "qasper.jsonl") - record = { - "dataset": "qasper", - "sample_idx": 0, - "source_idx": 12, - "status": "success", - "prompt_tokens": 16001, - "rendered_prompt": "fixed rendered prompt", - "rendered_prompt_sha256": longbench_pred._sha256_text( - "fixed rendered prompt" - ), - "pred": "answer", - "raw_pred": "answer", - "answers": ["answer"], - "all_classes": [], - "length": 16001, - } - - longbench_pred._write_sample_record( - out_root=tmp, - task_out_path=task_path, - record=record, - ) - - canonical = (Path(tmp) / "per_sample_results.jsonl").read_text( - encoding="utf-8" - ) - historical = (Path(tmp) / "sample_results.jsonl").read_text( - encoding="utf-8" - ) - raw = json.loads( - (Path(tmp) / "raw_outputs.jsonl").read_text(encoding="utf-8") - ) - self.assertEqual(canonical, historical) - self.assertEqual(raw["rendered_prompt"], "fixed rendered prompt") - self.assertEqual( - raw["rendered_prompt_sha256"], - longbench_pred._sha256_text("fixed rendered prompt"), - ) - def test_sparsevllm_data_workers_receive_distinct_master_ports(self): launched = [] @@ -137,27 +95,18 @@ def fake_popen(command, *, env, cwd): ) def test_longbench_records_actual_decode_cuda_graph_state(self): - worker_statuses = [ - { - "world_rank": 0, - "decode_cuda_graph_configured": True, - "decode_cuda_graph_state_count": 2, - "decode_cuda_graph_graph_count": 1, - "decode_cuda_graph_active": True, - }, - { - "world_rank": 1, - "decode_cuda_graph_configured": True, - "decode_cuda_graph_state_count": 2, - "decode_cuda_graph_graph_count": 1, - "decode_cuda_graph_active": True, + graph_runner = SimpleNamespace( + _graphs={ + "captured": SimpleNamespace(graph=object()), + "uncaptured": SimpleNamespace(graph=None), }, - ] + last_state_key="captured", + ) generate_fn = SimpleNamespace( _sparsevllm_llm=SimpleNamespace( - config=SimpleNamespace(world_size=2), + config=SimpleNamespace(decode_cuda_graph=True), model_runner=SimpleNamespace( - call=lambda method: worker_statuses + decode_cuda_graph_runner=graph_runner, ), ) ) @@ -170,71 +119,18 @@ def test_longbench_records_actual_decode_cuda_graph_state(self): ) path = Path(tmp) / "decode_cuda_graph_status_rank2.json" - self.assertEqual(status["launcher_rank"], 2) - self.assertTrue(status["configured_on_all_workers"]) - self.assertTrue(status["active_on_all_workers"]) - self.assertEqual(status["workers"], worker_statuses) + self.assertEqual(status["rank"], 2) + self.assertTrue(status["configured"]) + self.assertTrue(status["runner_initialized"]) + self.assertEqual(status["state_count"], 2) + self.assertEqual(status["graph_count"], 1) + self.assertTrue(status["active"]) + self.assertEqual(status["last_state_key"], "captured") self.assertEqual( json.loads(path.read_text(encoding="utf-8")), status, ) - def test_longbench_fails_if_any_requested_graph_worker_is_inactive(self): - generate_fn = SimpleNamespace( - _sparsevllm_llm=SimpleNamespace( - config=SimpleNamespace(world_size=2), - model_runner=SimpleNamespace( - call=lambda method: [ - { - "world_rank": 0, - "decode_cuda_graph_configured": True, - "decode_cuda_graph_active": True, - }, - { - "world_rank": 1, - "decode_cuda_graph_configured": True, - "decode_cuda_graph_active": False, - }, - ] - ), - ) - ) - - with tempfile.TemporaryDirectory() as tmp: - with self.assertRaisesRegex(RuntimeError, "not active on every"): - longbench_pred._write_decode_cuda_graph_status( - generate_fn=generate_fn, - out_root=tmp, - rank=0, - ) - - def test_longbench_fails_if_requested_graph_is_unconfigured(self): - generate_fn = SimpleNamespace( - _sparsevllm_llm=SimpleNamespace( - config=SimpleNamespace( - world_size=1, - decode_cuda_graph=True, - ), - model_runner=SimpleNamespace( - call=lambda method: [ - { - "world_rank": 0, - "decode_cuda_graph_configured": False, - "decode_cuda_graph_active": False, - } - ] - ), - ) - ) - - with tempfile.TemporaryDirectory() as tmp: - with self.assertRaisesRegex(RuntimeError, "not configured on every"): - longbench_pred._write_decode_cuda_graph_status( - generate_fn=generate_fn, - out_root=tmp, - rank=0, - ) - def test_longbench_fails_if_sparsevllm_graph_state_is_unavailable(self): with tempfile.TemporaryDirectory() as tmp: with self.assertRaisesRegex(RuntimeError, "_sparsevllm_llm"): diff --git a/tests/test_microbench_artifacts.py b/tests/test_microbench_artifacts.py index 3f4f2b7d..3bdc0c99 100644 --- a/tests/test_microbench_artifacts.py +++ b/tests/test_microbench_artifacts.py @@ -9,32 +9,10 @@ _benchmark_sparse_method, _record_child_exit_failure, _resolved_engine_config, - _worker_runtime_status, _write_output_dir, ) -def test_worker_runtime_status_preserves_per_rank_graph_and_provider_data(): - expected = [ - { - "world_rank": 0, - "decode_cuda_graph_active": True, - "moe_router_provider": "triton", - }, - { - "world_rank": 1, - "decode_cuda_graph_active": True, - "moe_router_provider": "triton", - }, - ] - llm = SimpleNamespace( - config=SimpleNamespace(world_size=2), - model_runner=SimpleNamespace(call=lambda method: expected), - ) - - assert _worker_runtime_status(llm) == expected - - @pytest.mark.parametrize( ("method", "expected"), [ diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 81ca023b..6e05b2e4 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -1,4 +1,3 @@ -import os import sys from types import SimpleNamespace from unittest.mock import Mock, patch @@ -13,16 +12,7 @@ TritonFp8LinearProvider, resolve_fp8_linear_provider, ) -from sparsevllm.operators.moe import ( - MOE_REGISTRY, - MoeOpSpec, - TorchMoeProvider, - resolve_moe_provider, -) -from sparsevllm.operators.moe_router import ( - MoeRouterOpSpec, - resolve_moe_router_provider, -) +from sparsevllm.operators.moe import MOE_REGISTRY, MoeOpSpec, resolve_moe_provider from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum @@ -74,7 +64,6 @@ def _moe_spec( tp_size=1, routing_method="softmax", scale_dtype=None, - cuda_graph=True, ) -> MoeOpSpec: return MoeOpSpec( num_experts=num_experts, @@ -86,7 +75,7 @@ def _moe_spec( weight_dtype=weight_dtype, block_shape=block_shape, ep_size=ep_size, - cuda_graph=cuda_graph, + cuda_graph=True, tp_size=tp_size, routing_method=routing_method, scale_dtype=scale_dtype, @@ -357,110 +346,6 @@ def test_unquantized_moe_uses_triton_on_supported_cuda(dtype, capability): assert resolved.provider.name == "triton" -def test_torch_moe_reference_matches_explicit_expert_routing(): - spec = _moe_spec( - activation_dtype=torch.float32, - weight_dtype=torch.float32, - block_shape=None, - hidden_size=3, - intermediate_size=2, - num_local_experts=2, - num_experts=4, - top_k=2, - ep_size=2, - cuda_graph=False, - ) - hidden_states = torch.tensor( - [[1.0, -2.0, 0.5], [0.25, 1.0, -0.75], [-1.0, 0.5, 2.0]] - ) - topk_ids = torch.tensor([[2, 0], [3, 2], [1, 0]]) - topk_weights = torch.tensor([[0.7, 0.3], [0.4, 0.6], [0.8, 0.2]]) - w13_weight = torch.arange(24, dtype=torch.float32).reshape(2, 4, 3) / 17 - w2_weight = torch.arange(12, dtype=torch.float32).reshape(2, 3, 2) / 11 - - actual = TorchMoeProvider().run( - spec, - hidden_states, - topk_ids, - topk_weights, - w13_weight, - w2_weight, - None, - None, - local_expert_start=2, - ep_rank=1, - ) - - expected = torch.zeros_like(hidden_states) - for token_index, routes in enumerate(zip(topk_ids, topk_weights)): - for expert_id, route_weight in zip(*routes): - local_expert_id = int(expert_id) - 2 - if not 0 <= local_expert_id < 2: - continue - gate_up = hidden_states[token_index] @ w13_weight[local_expert_id].T - gate, up = gate_up.chunk(2) - expert_output = ( - torch.nn.functional.silu(gate) * up - ) @ w2_weight[local_expert_id].T - expected[token_index] += expert_output * route_weight - - torch.testing.assert_close(actual, expected) - assert torch.equal(actual[2], torch.zeros(3)) - - -def test_explicit_torch_moe_provider_selection_is_eager_only(): - caps = _cuda_caps((9, 0), native_fp8=False) - platform = SimpleNamespace(get_device_caps=lambda _: caps) - eager_spec = _moe_spec( - activation_dtype=torch.bfloat16, - weight_dtype=torch.bfloat16, - block_shape=None, - cuda_graph=False, - ) - graph_spec = _moe_spec( - activation_dtype=torch.bfloat16, - weight_dtype=torch.bfloat16, - block_shape=None, - cuda_graph=True, - ) - - with ( - patch("sparsevllm.operators.moe.platforms.current_platform", platform), - patch.dict(os.environ, {"SPARSEVLLM_MOE_PROVIDER": "torch"}), - ): - assert resolve_moe_provider(eager_spec, device_index=0).name == "torch" - with pytest.raises(RuntimeError, match="eager-only"): - resolve_moe_provider(graph_spec, device_index=0) - - -def test_explicit_torch_router_provider_selection_is_eager_only(): - caps = _cuda_caps((9, 0), native_fp8=False) - platform = SimpleNamespace(get_device_caps=lambda _: caps) - - def spec(cuda_graph): - return MoeRouterOpSpec( - num_experts=256, - top_k=8, - activation_dtype=torch.bfloat16, - norm_topk_prob=True, - cuda_graph=cuda_graph, - ) - - with ( - patch( - "sparsevllm.operators.moe_router.platforms.current_platform", - platform, - ), - patch.dict( - os.environ, - {"SPARSEVLLM_MOE_ROUTER_PROVIDER": "torch"}, - ), - ): - assert resolve_moe_router_provider(spec(False), device_index=0).name == "torch" - with pytest.raises(RuntimeError, match="eager-only"): - resolve_moe_router_provider(spec(True), device_index=0) - - def test_hopper_fused_moe_uses_profiled_tp_ep_shape(): spec = _moe_spec( activation_dtype=torch.bfloat16, diff --git a/tests/test_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index 43b03bdf..f1f44396 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -1685,7 +1685,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", num_experts=128) + config.hf_config = SimpleNamespace(model_type="qwen3_moe") self.assertEqual( _moe_workspace_warmup_token_counts(config), @@ -1704,7 +1704,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", num_experts=128), + hf_config=SimpleNamespace(model_type="qwen3_moe"), ) calls = [] engine.model_runner = SimpleNamespace( @@ -1727,7 +1727,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", num_experts=128), + hf_config=SimpleNamespace(model_type="qwen3_moe"), ) def fail_on_workspace(_method, _num_tokens): diff --git a/tests/test_qwen35_mixed_runtime.py b/tests/test_qwen35_mixed_runtime.py index 5bc3a2c5..24ae8118 100644 --- a/tests/test_qwen35_mixed_runtime.py +++ b/tests/test_qwen35_mixed_runtime.py @@ -54,37 +54,6 @@ def _single_process_parallel_context() -> ParallelContext: return ParallelContext(world=group, tensor=group, expert=group, data=group) -def test_model_runner_debug_logits_capture_is_explicit(): - runner = SimpleNamespace(debug_runtime_enabled=True) - - ModelRunner._capture_debug_logits(runner, None) - assert not hasattr(runner, "debug_last_logits") - - logits = torch.tensor([[1.0, 2.0]]) - ModelRunner._capture_debug_logits(runner, logits) - torch.testing.assert_close(runner.debug_last_logits, logits) - assert runner.debug_last_logits.data_ptr() != logits.data_ptr() - - -def test_model_runner_debug_logits_non_output_rank_returns_none(): - runner = SimpleNamespace(rank=1) - - assert ModelRunner.debug_last_logits_cpu(runner) is None - - -def test_model_runner_tp_replica_consistency_does_not_require_rank_local_logits(): - runner = SimpleNamespace( - world_size=2, - model=SimpleNamespace(model=SimpleNamespace(layers=[])), - ) - - assert ModelRunner.debug_replica_consistency(runner) == { - "last_logits_max_abs": None, - "last_logits_tolerance_ratio": None, - "moe_layers": {}, - } - - def _qwen35_outer_config(*, num_layers: int = 64, full_layers: tuple[int, ...] | None = None): if full_layers is None: full_layers = tuple(range(0, num_layers, 4)) diff --git a/tests/test_qwen35_moe.py b/tests/test_qwen35_moe.py deleted file mode 100644 index 46ed1494..00000000 --- a/tests/test_qwen35_moe.py +++ /dev/null @@ -1,563 +0,0 @@ -from contextlib import ExitStack -from types import SimpleNamespace -from unittest.mock import Mock, patch - -import pytest -import torch - -from sparsevllm.config import Config, QuantizationConfig -from sparsevllm.distributed import ParallelContext, ParallelGroup -from sparsevllm.engine.recurrent_state_manager import RecurrentStateManager -from sparsevllm.models.qwen3_5_moe import ( - Qwen35MoeForCausalLM, - Qwen35MoePackedExperts, - Qwen35MoeSparseMoeBlock, -) -from sparsevllm.models.qwen3_5 import Qwen35LinearAttention -from sparsevllm.operators.moe import TritonMoeProvider -from sparsevllm.operators.moe_router import ( - MoeRouterOpSpec, - TorchMoeRouterProvider, -) - - -def _outer_config(*, fp8: bool = False): - layer_types = [ - "full_attention" if (layer_idx + 1) % 4 == 0 else "linear_attention" - for layer_idx in range(40) - ] - text_config = SimpleNamespace( - model_type="qwen3_5_moe_text", - vocab_size=248320, - hidden_size=2048, - num_hidden_layers=40, - layer_types=layer_types, - num_attention_heads=16, - num_key_value_heads=2, - head_dim=256, - linear_num_key_heads=16, - linear_num_value_heads=32, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_conv_kernel_dim=4, - moe_intermediate_size=512, - shared_expert_intermediate_size=512, - num_experts=256, - num_experts_per_tok=8, - hidden_act="silu", - attn_output_gate=True, - attention_bias=False, - partial_rotary_factor=0.25, - mamba_ssm_dtype="float32", - tie_word_embeddings=False, - rms_norm_eps=1.0e-6, - max_position_embeddings=262144, - torch_dtype=torch.bfloat16, - quantization_config=None, - ) - outer_config = SimpleNamespace( - model_type="qwen3_5_moe", - architectures=["Qwen3_5MoeForConditionalGeneration"], - text_config=text_config, - ) - if fp8: - del text_config.quantization_config - outer_config.quantization_config = { - "quant_method": "fp8", - "fmt": "e4m3", - "activation_scheme": "dynamic", - "weight_block_size": [128, 128], - } - return outer_config - - -def _make_config(tmp_path, **overrides): - with patch( - "sparsevllm.configs.runtime.AutoConfig.from_pretrained", - return_value=_outer_config(), - ): - return Config(model=str(tmp_path), **overrides) - - -def _hybrid_context(world_rank: int) -> ParallelContext: - world_ranks = (0, 1, 2, 3) - moe_tp_ranks = (0, 1) if world_rank < 2 else (2, 3) - moe_ep_ranks = (0, 2) if world_rank % 2 == 0 else (1, 3) - return ParallelContext( - world=ParallelGroup(None, world_ranks, world_rank, 4), - tensor=ParallelGroup(None, world_ranks, world_rank, 4), - expert=ParallelGroup( - None, moe_ep_ranks, moe_ep_ranks.index(world_rank), 2 - ), - data=ParallelGroup(None, (world_rank,), 0, 1), - moe_tensor=ParallelGroup( - None, moe_tp_ranks, moe_tp_ranks.index(world_rank), 2 - ), - ) - - -def _single_context() -> ParallelContext: - group = ParallelGroup(None, (0,), 0, 1) - return ParallelContext(group, group, group, group) - - -def _pure_tp_context(world_rank: int) -> ParallelContext: - world = ParallelGroup(None, (0, 1), world_rank, 2) - singleton = ParallelGroup(None, (world_rank,), 0, 1) - return ParallelContext( - world=world, - tensor=world, - expert=singleton, - data=singleton, - moe_tensor=world, - ) - - -def _fp8_expert_config(): - return SimpleNamespace( - num_experts=2, - hidden_size=128, - moe_intermediate_size=128, - num_experts_per_tok=1, - dtype=torch.bfloat16, - torch_dtype=torch.bfloat16, - decode_cuda_graph=True, - quantization_config=QuantizationConfig( - enabled=True, - quant_method="fp8", - weight_dtype="e4m3", - activation_scheme="dynamic", - weight_block_size=(128, 128), - model_name="Qwen3.6 MoE", - ), - ) - - -def _make_fp8_experts(): - with ( - patch( - "sparsevllm.models.qwen3_moe.get_parallel_context", - return_value=_single_context(), - ), - patch( - "sparsevllm.models.qwen3_moe.resolve_moe_provider", - return_value=TritonMoeProvider(), - ), - ): - return Qwen35MoePackedExperts(_fp8_expert_config()) - - -def test_qwen36_moe_config_normalizes_text_runtime_and_topology(tmp_path): - config = _make_config( - tmp_path, - tensor_parallel_size=2, - expert_parallel_size=2, - decode_cuda_graph=True, - enforce_eager=False, - ) - - assert config.hf_config.model_type == "qwen3_5_moe" - assert config.uses_outer_tp_moe_layout is True - assert config.world_size == 2 - assert config.moe_tensor_parallel_size == 1 - assert config.runtime_layout.full_attention_layer_indices == tuple( - range(3, 40, 4) - ) - assert config.runtime_layout.num_kv_layers == 10 - - -@pytest.mark.parametrize( - "method", - [ - "streamingllm", - "snapkv", - "h2o", - "pyramidkv", - "omnikv", - "quest", - "rkv", - ], -) -def test_qwen36_moe_accepts_asset_free_sparse_graph_methods(tmp_path, method): - config = _make_config( - tmp_path, - vllm_sparse_method=method, - full_attn_layers="3,11,19,27,35", - decode_cuda_graph=True, - enforce_eager=False, - ) - - assert config.vllm_sparse_method == method - assert config.decode_cuda_graph is True - - -@pytest.mark.parametrize( - ("method", "error"), - [ - ("skipkv", "official models with released steering vectors"), - ("deltakv", "validated methods"), - ], -) -def test_qwen36_moe_rejects_sparse_methods_requiring_model_assets( - tmp_path, - method, - error, -): - with pytest.raises(ValueError, match=error): - _make_config(tmp_path, vllm_sparse_method=method) - - -def test_qwen36_moe_rejects_invalid_outer_tp_ep_topology(tmp_path): - with pytest.raises(ValueError, match="must be divisible"): - _make_config( - tmp_path, - tensor_parallel_size=2, - expert_parallel_size=3, - ) - - -def test_qwen36_moe_rejects_non_bf16_checkpoint(tmp_path): - outer = _outer_config() - outer.text_config.torch_dtype = torch.float16 - with patch( - "sparsevllm.configs.runtime.AutoConfig.from_pretrained", - return_value=outer, - ): - with pytest.raises(NotImplementedError, match="requires BF16"): - Config(model=str(tmp_path)) - - -def test_qwen36_moe_accepts_outer_block_fp8_config(tmp_path): - with patch( - "sparsevllm.configs.runtime.AutoConfig.from_pretrained", - return_value=_outer_config(fp8=True), - ): - config = Config(model=str(tmp_path)) - - assert config.quantization_config.enabled is True - assert config.quantization_config.weight_dtype == "e4m3" - assert config.quantization_config.weight_block_size == (128, 128) - - -def test_qwen36_moe_fp8_rejects_unsupported_outer_tp(tmp_path): - with patch( - "sparsevllm.configs.runtime.AutoConfig.from_pretrained", - return_value=_outer_config(fp8=True), - ): - with pytest.raises(ValueError, match="num_key_value_heads"): - Config(model=str(tmp_path), tensor_parallel_size=8) - - -def test_qwen36_moe_recurrent_state_uses_attention_tp_and_fp32_state(): - config = _outer_config().text_config - - spec = Qwen35MoeForCausalLM.recurrent_state_spec( - config, attention_tp_size=2 - ) - - assert spec.tensor_specs[0].shape == (4096, 3) - assert spec.tensor_specs[0].dtype == torch.bfloat16 - assert spec.tensor_specs[1].shape == (16, 128, 128) - assert spec.tensor_specs[1].dtype == torch.float32 - - -def test_qwen36_moe_torch_router_is_fp32_softmax_oracle(): - logits = torch.tensor( - [[-80.0, -2.0, 0.0, 1.0, 3.0, 7.0, 8.0, 9.0, 10.0]], - dtype=torch.bfloat16, - ) - spec = MoeRouterOpSpec( - num_experts=9, - top_k=8, - activation_dtype=torch.bfloat16, - norm_topk_prob=True, - cuda_graph=False, - ) - - weights, ids = TorchMoeRouterProvider().run(spec, logits) - probabilities = torch.softmax(logits, dim=-1, dtype=torch.float32) - expected_weights, expected_ids = torch.topk(probabilities, 8, dim=-1) - expected_weights /= expected_weights.sum(dim=-1, keepdim=True) - - assert weights.dtype == torch.bfloat16 - assert ids.dtype == torch.int32 - assert torch.equal(ids, expected_ids.to(torch.int32)) - torch.testing.assert_close( - weights.float(), expected_weights, atol=4e-3, rtol=4e-3 - ) - - -def test_qwen36_moe_linear_attention_uses_configured_recurrent_dtype(): - config = SimpleNamespace( - hidden_size=8, - hidden_act="silu", - linear_num_key_heads=1, - linear_num_value_heads=2, - linear_key_head_dim=4, - linear_value_head_dim=4, - linear_conv_kernel_dim=4, - rms_norm_eps=1.0e-6, - mlp_chunk_size=16, - torch_dtype=torch.bfloat16, - runtime_recurrent_state_dtype=torch.float32, - quantization_config=QuantizationConfig.disabled(), - ) - context = _single_context() - with ( - patch( - "sparsevllm.models.qwen3_5.get_parallel_context", - return_value=context, - ), - patch( - "sparsevllm.layers.linear.get_parallel_context", - return_value=context, - ), - ): - attention = Qwen35LinearAttention(config) - - assert attention.recurrent_state_dtype == torch.float32 - - -def test_recurrent_pool_accepts_model_declared_mixed_state_dtypes(): - runtime_config = SimpleNamespace( - runtime_layout=SimpleNamespace( - linear_attention_layer_indices=(0,), - is_linear_attention=lambda layer_idx: int(layer_idx) == 0, - ), - enable_prefix_caching=False, - max_num_seqs_in_batch=1, - max_decoding_seqs=1, - max_num_seqs_in_gpu=1, - recurrent_state_max_bytes=None, - prefix_cache_block_size=4, - ) - state_spec = Qwen35MoeForCausalLM.recurrent_state_spec( - _outer_config().text_config, - attention_tp_size=1, - ) - manager = RecurrentStateManager( - runtime_config, - _single_context(), - device=torch.device("cpu"), - state_spec=state_spec, - ) - seq = SimpleNamespace(seq_id=1) - manager.prepare_step([seq], is_prefill=False) - manager.prepare_decode_static([seq], token_batch=1, device=torch.device("cpu")) - - state_buffers, _ = manager.get_decode_layer_state( - [seq], - layer_idx=0, - token_batch=1, - dtype=torch.bfloat16, - device=torch.device("cpu"), - ) - - assert state_buffers["conv_state"].dtype == torch.bfloat16 - assert state_buffers["recurrent_state"].dtype == torch.float32 - - -def test_packed_experts_slice_ep_before_moe_tp(): - context = _hybrid_context(world_rank=1) - config = SimpleNamespace( - num_experts=4, - hidden_size=4, - moe_intermediate_size=4, - num_experts_per_tok=2, - dtype=torch.bfloat16, - torch_dtype=torch.bfloat16, - decode_cuda_graph=False, - quantization_config=QuantizationConfig.disabled(), - ) - previous_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.bfloat16) - try: - with ExitStack() as stack: - stack.enter_context( - patch( - "sparsevllm.models.qwen3_moe.get_parallel_context", - return_value=context, - ) - ) - stack.enter_context( - patch( - "sparsevllm.models.qwen3_moe.resolve_moe_provider", - return_value=TritonMoeProvider(), - ) - ) - experts = Qwen35MoePackedExperts(config) - finally: - torch.set_default_dtype(previous_dtype) - - gate_up_global_shape = (4, 8, 4) - down_global_shape = (4, 4, 4) - assert experts.rank_local_weight_slice( - gate_up_global_shape, - loaded_shard_id="gate_up_proj", - ) == (slice(0, 2), slice(None), slice(None)) - assert experts.rank_local_weight_slice( - down_global_shape, - loaded_shard_id="down_proj", - ) == (slice(0, 2), slice(None), slice(None)) - - gate_up = torch.arange(2 * 8 * 4, dtype=torch.bfloat16).view(2, 8, 4) - down = torch.arange(2 * 4 * 4, dtype=torch.bfloat16).view(2, 4, 4) - experts.load_packed_expert_weight("gate_up_proj", gate_up) - experts.load_packed_expert_weight("down_proj", down) - experts.validate_loaded_weights() - - expected_gate = gate_up[:, 2:4] - expected_up = gate_up[:, 6:8] - torch.testing.assert_close(experts.w13_weight[:, :2], expected_gate) - torch.testing.assert_close(experts.w13_weight[:, 2:], expected_up) - torch.testing.assert_close(experts.w2_weight, down[:, :, 2:4]) - - -def test_packed_expert_pure_tp_shards_reconstruct_checkpoint(): - config = SimpleNamespace( - num_experts=4, - hidden_size=4, - moe_intermediate_size=4, - num_experts_per_tok=2, - dtype=torch.bfloat16, - torch_dtype=torch.bfloat16, - decode_cuda_graph=False, - quantization_config=QuantizationConfig.disabled(), - ) - gate_up = torch.arange(4 * 8 * 4, dtype=torch.bfloat16).view(4, 8, 4) - down = torch.arange(4 * 4 * 4, dtype=torch.bfloat16).view(4, 4, 4) - rank_experts = [] - previous_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.bfloat16) - try: - for world_rank in (0, 1): - with ExitStack() as stack: - stack.enter_context( - patch( - "sparsevllm.models.qwen3_moe.get_parallel_context", - return_value=_pure_tp_context(world_rank), - ) - ) - stack.enter_context( - patch( - "sparsevllm.models.qwen3_moe.resolve_moe_provider", - return_value=TritonMoeProvider(), - ) - ) - experts = Qwen35MoePackedExperts(config) - experts.load_packed_expert_weight("gate_up_proj", gate_up) - experts.load_packed_expert_weight("down_proj", down) - experts.validate_loaded_weights() - rank_experts.append(experts) - finally: - torch.set_default_dtype(previous_dtype) - - gate = torch.cat( - [experts.w13_weight[:, :2] for experts in rank_experts], dim=1 - ) - up = torch.cat( - [experts.w13_weight[:, 2:] for experts in rank_experts], dim=1 - ) - reconstructed_down = torch.cat( - [experts.w2_weight for experts in rank_experts], dim=2 - ) - torch.testing.assert_close(gate, gate_up[:, :4]) - torch.testing.assert_close(up, gate_up[:, 4:]) - torch.testing.assert_close(reconstructed_down, down) - - -def test_qwen36_fp8_experts_load_per_expert_weights_and_scales(): - experts = _make_fp8_experts() - sources = {} - for expert_id in range(experts.num_experts): - for projection in ("gate_proj", "up_proj", "down_proj"): - weight = ( - torch.randn(128, 128) - .clamp(-4.0, 4.0) - .to(torch.float8_e4m3fn) - ) - scale = torch.rand(1, 1, dtype=torch.bfloat16) + 0.1 - sources[(expert_id, projection)] = (weight, scale) - experts.load_expert_weight(expert_id, projection, weight, scale) - - experts.validate_loaded_weights() - for expert_id in range(experts.num_experts): - gate, gate_scale = sources[(expert_id, "gate_proj")] - up, up_scale = sources[(expert_id, "up_proj")] - down, down_scale = sources[(expert_id, "down_proj")] - assert torch.equal(experts.w13_weight[expert_id, :128], gate) - assert torch.equal(experts.w13_weight[expert_id, 128:], up) - assert torch.equal(experts.w13_scale_inv[expert_id, :1], gate_scale) - assert torch.equal(experts.w13_scale_inv[expert_id, 1:], up_scale) - assert torch.equal(experts.w2_weight[expert_id], down) - assert torch.equal(experts.w2_scale_inv[expert_id], down_scale) - - -def test_qwen36_checkpoint_adapter_keeps_fp8_layout_model_local(): - experts = _make_fp8_experts() - model = Qwen35MoeForCausalLM.__new__(Qwen35MoeForCausalLM) - torch.nn.Module.__init__(model) - model.model = torch.nn.Module() - model.model.layers = torch.nn.ModuleList([torch.nn.Module()]) - model.model.layers[0].mlp = torch.nn.Module() - model.model.layers[0].mlp.experts = experts - - source_name = ( - "model.language_model.layers.0.mlp.experts.1.gate_proj.weight" - ) - target_name = model.map_weight_name(source_name) - - assert target_name == "model.layers.0.mlp.experts.1.gate_proj.expert_weight" - target = model.resolve_special_weight(target_name) - assert target is not None - assert target.module is experts - assert target.shard_id == (1, "gate_proj") - with pytest.raises(ValueError, match="must use per-expert"): - model.map_weight_name( - "model.language_model.layers.0.mlp.experts.gate_up_proj" - ) - - -def test_qwen36_fp8_experts_require_scale(): - experts = _make_fp8_experts() - weight = torch.ones(128, 128).to(torch.float8_e4m3fn) - - with pytest.raises(ValueError, match="Missing FP8 weight_scale_inv"): - experts.load_expert_weight(0, "gate_proj", weight, None) - - -def test_routed_output_reduces_without_reducing_shared_output_twice(): - class FixedRouter(torch.nn.Module): - def forward(self, hidden_states): - tokens = hidden_states.shape[0] - return ( - torch.zeros(tokens, 4), - torch.full((tokens, 2), 0.5), - torch.zeros(tokens, 2, dtype=torch.int32), - ) - - class FixedExperts(torch.nn.Module): - def forward(self, hidden_states, _topk_ids, _topk_weights): - return torch.full_like(hidden_states, 2.0) - - class ZeroGate(torch.nn.Module): - def forward(self, hidden_states): - return torch.zeros(hidden_states.shape[0], 1) - - block = Qwen35MoeSparseMoeBlock.__new__(Qwen35MoeSparseMoeBlock) - torch.nn.Module.__init__(block) - block.shared_expert = torch.nn.Identity() - block.shared_expert_gate = ZeroGate() - block.gate = FixedRouter() - block.experts = FixedExperts() - block.parallel_context = SimpleNamespace( - world_all_reduce=Mock(side_effect=lambda tensor: tensor * 3) - ) - hidden_states = torch.full((2, 3), 4.0) - - output, *_ = block._forward_chunk(hidden_states) - - torch.testing.assert_close(output, torch.full_like(hidden_states, 8.0)) - block.parallel_context.world_all_reduce.assert_called_once() - reduced_input = block.parallel_context.world_all_reduce.call_args.args[0] - torch.testing.assert_close(reduced_input, torch.full_like(hidden_states, 2.0)) diff --git a/tests/test_qwen36_longbench_compare.py b/tests/test_qwen36_longbench_compare.py deleted file mode 100644 index 50ad1263..00000000 --- a/tests/test_qwen36_longbench_compare.py +++ /dev/null @@ -1,46 +0,0 @@ -from pathlib import Path - -import pytest - -from scripts.validation.qwen36_compare_longbench import ( - _validate_worker_providers, -) - - -def test_fp8_longbench_accepts_registered_graph_providers(): - workers = [ - { - "moe_expert_provider": "flashinfer_cutlass_fp8_sm90", - "moe_router_provider": "triton", - "moe_weight_dtype": "torch.float8_e4m3fn", - "fp8_linear_provider": "flashinfer_sm90", - }, - { - "moe_expert_provider": "triton", - "moe_router_provider": "triton", - "moe_weight_dtype": "torch.float8_e4m3fn", - "fp8_linear_provider": "flashinfer_sm90", - }, - ] - - _validate_worker_providers( - workers, - precision="fp8", - path=Path("fp8-run"), - ) - - -def test_fp8_longbench_rejects_missing_fp8_diagnostics(): - workers = [ - { - "moe_expert_provider": "triton", - "moe_router_provider": "triton", - } - ] - - with pytest.raises(RuntimeError, match="invalid FP8 providers"): - _validate_worker_providers( - workers, - precision="fp8", - path=Path("fp8-run"), - ) diff --git a/tests/test_sparsevllm_regression_grading.py b/tests/test_sparsevllm_regression_grading.py index 2a80f8a4..7e5a1d41 100644 --- a/tests/test_sparsevllm_regression_grading.py +++ b/tests/test_sparsevllm_regression_grading.py @@ -191,15 +191,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", - "qwen3_5_moe", - "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 141ea2c9..50842041 100644 --- a/tests/test_tp_rpc.py +++ b/tests/test_tp_rpc.py @@ -258,10 +258,6 @@ def test_moe_workspace_warmup_uses_failure_synchronized_world_rpc(): assert "warmup_moe_workspace" in TP_RPC_STATUS_SYNC_METHODS -def test_runtime_diagnostics_use_failure_synchronized_world_rpc(): - assert "runtime_diagnostic_status" in TP_RPC_STATUS_SYNC_METHODS - - def test_fake_prefill_warmup_uses_failure_synchronized_world_rpc(): assert "set_warmup_fake_prefill_attention" in TP_RPC_STATUS_SYNC_METHODS diff --git a/tests/test_triton_moe.py b/tests/test_triton_moe.py index 7e101e1f..acd90195 100644 --- a/tests/test_triton_moe.py +++ b/tests/test_triton_moe.py @@ -93,15 +93,14 @@ def test_moe_alignment_covers_hotspot_and_empty_rank(dtype): assert int(valid.numel()) == expected -@pytest.mark.parametrize("num_experts", [128, 256]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("norm_topk_prob", [False, True]) @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for Triton MoE tests.") -def test_topk_softmax_matches_pytorch(num_experts, dtype, norm_topk_prob): +def test_topk_softmax_matches_pytorch(dtype, norm_topk_prob): torch.manual_seed(21) - base = torch.arange(num_experts, dtype=dtype, device="cuda") / 16 - 4 + base = torch.arange(128, dtype=dtype, device="cuda") / 16 - 4 logits = torch.stack( - [base[torch.randperm(num_experts, device="cuda")] for _ in range(257)] + [base[torch.randperm(128, device="cuda")] for _ in range(257)] ) expected_probs = torch.softmax(logits, dim=-1, dtype=torch.float32) expected_weights, expected_ids = torch.topk(expected_probs, 8, dim=-1) @@ -167,7 +166,7 @@ def test_topk_softmax_nonfinite_inputs_keep_ids_in_range_and_propagate_nan(): @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for Triton MoE tests.") def test_topk_softmax_rejects_unsupported_shape_and_layout(): - with pytest.raises(ValueError, match="num_experts in"): + with pytest.raises(ValueError, match="num_experts=128"): topk_softmax( torch.zeros(2, 64, dtype=torch.bfloat16, device="cuda"), top_k=8, From a23d354512baa477fdd64684687c2b89a63a480d Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 14:54:40 +0800 Subject: [PATCH 05/35] fix: update Qwen3.6 MoE support details in documentation --- README.md | 2 +- README_zh.md | 2 +- docs/en/features/supported-models.md | 11 +---------- docs/zh/features/supported-models.md | 10 +--------- 4 files changed, 4 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 0ed85ef7..5b28e54a 100644 --- a/README.md +++ b/README.md @@ -69,7 +69,7 @@ Read the method overview and integration rules in | Qwen3 | ✅ | | Qwen3MoE | ✅ | | Qwen3.5 / Qwen3.6 | ✅ | -| Qwen3.6 MoE | ✅ (BF16 / block FP8, text-only) | +| Qwen3.5 / Qwen3.6 MoE | ✅ | | Llama 3 / 3.1 | ✅ | | MiniMax M2.7 | ✅ | diff --git a/README_zh.md b/README_zh.md index c18da47f..b8351cfc 100644 --- a/README_zh.md +++ b/README_zh.md @@ -58,7 +58,7 @@ Sparse-vLLM 支持物理淘汰、逻辑掩码、查询感知选择和混合 KV | Qwen3 | ✅ | | Qwen3MoE | ✅ | | Qwen3.5 / Qwen3.6 | ✅ | -| Qwen3.6 MoE | ✅（仅 BF16 纯文本） | +| Qwen3.5 / Qwen3.6 MoE | ✅ | | Llama 3 / 3.1 | ✅ | | MiniMax M2.7 | ✅ | diff --git a/docs/en/features/supported-models.md b/docs/en/features/supported-models.md index d9126a97..8ff710a8 100644 --- a/docs/en/features/supported-models.md +++ b/docs/en/features/supported-models.md @@ -16,7 +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/2 | 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,15 +32,6 @@ dimension must be divisible by `T / E`. Qwen3MoE outer TP requires a BF16 model dtype; FP16 Qwen3MoE checkpoints are limited to `TP=1`. When `TP=1`, the existing EP layout uses world size `E`. -Qwen3.6 MoE always uses the outer-TP layout: attention and Gated DeltaNet TP -are `T`, MoE EP is `E`, MoE TP is `T / E`, and world size is `T`. It requires -`DP=1`, `T % E == 0`, and BF16 activations with either BF16 or block FP8 -language-model weights. The runtime is text-only, rejects image/video and MTP -inputs, and captures decode (not prefill) with CUDA Graph. Sparse methods apply -only to the full-attention layers; Gated DeltaNet layers keep their recurrent -state path. Outer TP is limited to 1 or 2 by the two KV heads; FP8 also requires -every TP-local quantized Linear dimension to remain 128-aligned. - Block FP8 support requires E4M3 weights, dynamic activation quantization, and 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 diff --git a/docs/zh/features/supported-models.md b/docs/zh/features/supported-models.md index 5ba0f170..44331922 100644 --- a/docs/zh/features/supported-models.md +++ b/docs/zh/features/supported-models.md @@ -12,7 +12,7 @@ | 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/2 | 仅支持 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 | ✅ | @@ -25,14 +25,6 @@ size 为 `T`。该布局要求 `DP=1` 且 `T % E == 0`；专家数量必须能 TP 要求模型 dtype 为 BF16；FP16 Qwen3MoE checkpoint 仅支持 `TP=1`。当 `TP=1` 时，原有 EP 布局的 world size 为 `E`。 -Qwen3.6 MoE 始终使用 outer-TP 布局：attention 与 Gated DeltaNet TP 为 -`T`、MoE EP 为 `E`、MoE TP 为 `T / E`，world size 为 `T`。该模型要求 -`DP=1`、`T % E == 0`，激活为 BF16，语言模型权重可使用 BF16 或块级 FP8。 -当前 runtime 仅支持纯文本 CausalLM，明确拒绝 image/video 与 MTP 输入。 -两个 KV heads 将 outer TP 限定为 1 或 2；CUDA Graph 仅覆盖 decode，不覆盖 -prefill。稀疏方法只作用于 full-attention 层，Gated DeltaNet 层继续使用独立 -的递归状态路径；FP8 还要求所有 TP-local 量化 Linear 维度保持 128 对齐。 - 块级 FP8 要求使用 E4M3 权重、动态激活量化以及 `128 x 128` 的权重块大小。 Qwen3.5/Qwen3.6 Dense 配置在内部统一规范为 `model_type=qwen3_5`；Qwen3.6 MoE 使用 `model_type=qwen3_5_moe`。 From fa227b50542e3f0a55b66afcdace1a635539af65 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 15:42:51 +0800 Subject: [PATCH 06/35] feat: log selected operator implementations --- src/sparsevllm/engine/llm_engine.py | 1 + src/sparsevllm/engine/model_runner.py | 5 +++ src/sparsevllm/layers/attention_backend.py | 6 ++++ src/sparsevllm/layers/layernorm.py | 3 ++ src/sparsevllm/operators/fp8_linear.py | 4 +++ src/sparsevllm/operators/registry.py | 29 +++++++++++++++ tests/test_operator_providers.py | 2 ++ tests/test_operator_registry.py | 41 ++++++++++++++++++++-- tests/test_prefill_schedule_policy.py | 1 + tests/test_tp_rpc.py | 12 +++++++ 10 files changed, 102 insertions(+), 2 deletions(-) diff --git a/src/sparsevllm/engine/llm_engine.py b/src/sparsevllm/engine/llm_engine.py index 86cf386a..2f7515d3 100644 --- a/src/sparsevllm/engine/llm_engine.py +++ b/src/sparsevllm/engine/llm_engine.py @@ -458,6 +458,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: diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 4fe06145..c81c4b2a 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 @@ -85,6 +86,7 @@ "debug_moe_states_cpu", "free_slots", "free_slots_batch", + "log_operator_implementations", "refresh_prefix_cache_hit", "reset_after_warmup", "run", @@ -586,6 +588,9 @@ 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: + operator_registry.log_operator_implementations(self.parallel_context.world_rank) + def warmup_moe_workspace(self, num_tokens: int) -> None: warmup_moe = getattr(self.model, "warmup_moe", None) if not callable(warmup_moe): 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/operators/fp8_linear.py b/src/sparsevllm/operators/fp8_linear.py index 068d9ca6..33d43d7a 100644 --- a/src/sparsevllm/operators/fp8_linear.py +++ b/src/sparsevllm/operators/fp8_linear.py @@ -56,6 +56,10 @@ class FlashInferSm90Fp8LinearProvider(Fp8LinearProvider): def __init__(self) -> None: self._fallback: TritonFp8LinearProvider | None = None + @property + def implementation_name(self) -> str: + return self._fallback.name if self._fallback is not None else self.name + @classmethod def supports(cls, spec: Fp8LinearSpec, caps: DeviceCaps) -> SupportResult: if spec.block_shape != (128, 128): diff --git a/src/sparsevllm/operators/registry.py b/src/sparsevllm/operators/registry.py index e174d301..bba92ccf 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(world_rank: int) -> 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 (rank {}):\n{}", world_rank, 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/tests/test_operator_providers.py b/tests/test_operator_providers.py index 6e05b2e4..d53479f6 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -199,6 +199,7 @@ def test_flashinfer_linear_binds_triton_for_missing_uncached_kernel(): ) fallback_output = torch.ones(2, 128, dtype=torch.bfloat16) provider = FlashInferSm90Fp8LinearProvider() + assert provider.implementation_name == "flashinfer_sm90" x = torch.ones(2, 128, dtype=torch.bfloat16) weight = torch.ones(128, 128).to(torch.float8_e4m3fn) scale = torch.ones(1, 1) @@ -226,6 +227,7 @@ def test_flashinfer_linear_binds_triton_for_missing_uncached_kernel(): assert second is fallback_output assert flashinfer_call.call_count == 1 assert triton_call.call_count == 2 + assert provider.implementation_name == "triton" def test_flashinfer_linear_does_not_mask_other_runtime_failures(): diff --git a/tests/test_operator_registry.py b/tests/test_operator_registry.py index d3cd4573..c4dddd7c 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,33 @@ 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(3) + + log_info.assert_called_once_with( + "Operator implementations (rank {}):\n{}", + 3, + " 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_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index f1f44396..9fba25b8 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -1639,6 +1639,7 @@ def step(): [ ("set_warmup_fake_prefill_attention", True, 2046), ("set_warmup_fake_prefill_attention", False), + ("log_operator_implementations",), ], ) diff --git a/tests/test_tp_rpc.py b/tests/test_tp_rpc.py index 50842041..15c27937 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,17 @@ 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(): + runner = object.__new__(ModelRunner) + runner.parallel_context = SimpleNamespace(world_rank=2) + + with patch.object(operator_registry, "log_operator_implementations") as log_implementations: + ModelRunner.log_operator_implementations(runner) + + assert "log_operator_implementations" in TP_RPC_STATUS_SYNC_METHODS + log_implementations.assert_called_once_with(2) + + 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) From 26da928419fb165fbd355439541bbaf3e303a8ed Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 15:52:51 +0800 Subject: [PATCH 07/35] feat: enable qwen3.6 moe prefix cache --- src/sparsevllm/method_registry.py | 2 +- tests/test_qwen3_moe_compatibility.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/src/sparsevllm/method_registry.py b/src/sparsevllm/method_registry.py index d95b0af8..07e81a19 100644 --- a/src/sparsevllm/method_registry.py +++ b/src/sparsevllm/method_registry.py @@ -145,7 +145,7 @@ class ModelRuntimeCompatibility: QWEN35_MOE_COMPATIBILITY = ModelRuntimeCompatibility( parallel_mode="outer_tp_moe_tp_ep", sparse_methods=QWEN3_MOE_TP_EP_COMPATIBILITY.sparse_methods, - prefix_cache_methods=frozenset(), + prefix_cache_methods=frozenset({""}), requires_eager=False, decode_cuda_graph_methods=( QWEN3_MOE_TP_EP_COMPATIBILITY.decode_cuda_graph_methods diff --git a/tests/test_qwen3_moe_compatibility.py b/tests/test_qwen3_moe_compatibility.py index 47aa506c..5b8bd379 100644 --- a/tests/test_qwen3_moe_compatibility.py +++ b/tests/test_qwen3_moe_compatibility.py @@ -2,6 +2,7 @@ from sparsevllm.method_registry import ( MODEL_RUNTIME_COMPATIBILITY, + QWEN35_MOE_COMPATIBILITY, QWEN3_MOE_EP_COMPATIBILITY, QWEN3_MOE_TP_COMPATIBILITY, QWEN3_MOE_TP_EP_COMPATIBILITY, @@ -59,6 +60,19 @@ 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="", + tensor_parallel_size=2, + expert_parallel_size=2, + data_parallel_size=1, + 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 From 9a01521e34df2251f572d1b91d0d8e0d6ae8745a Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 16:13:59 +0800 Subject: [PATCH 08/35] refactor: log operator implementations once --- src/sparsevllm/engine/model_runner.py | 3 ++- src/sparsevllm/operators/registry.py | 4 ++-- tests/test_operator_registry.py | 5 ++--- tests/test_tp_rpc.py | 11 +++++++---- 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index c81c4b2a..fd481639 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -589,7 +589,8 @@ def reset_after_warmup(self) -> None: self.sparse_controller.clear_decode_attn_score_buffers() def log_operator_implementations(self) -> None: - operator_registry.log_operator_implementations(self.parallel_context.world_rank) + 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) diff --git a/src/sparsevllm/operators/registry.py b/src/sparsevllm/operators/registry.py index bba92ccf..2081ac7c 100644 --- a/src/sparsevllm/operators/registry.py +++ b/src/sparsevllm/operators/registry.py @@ -23,7 +23,7 @@ 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(world_rank: int) -> None: +def log_operator_implementations() -> None: entries = sorted( ( operator_type, @@ -35,7 +35,7 @@ def log_operator_implementations(world_rank: int) -> None: if not entries: return rows = "\n".join(f" {operator_type}: {implementation}" for operator_type, implementation in entries) - logger.info("Operator implementations (rank {}):\n{}", world_rank, rows) + logger.info("Operator implementations:\n{}", rows) def runtime_version_at_least( diff --git a/tests/test_operator_registry.py b/tests/test_operator_registry.py index c4dddd7c..2070f828 100644 --- a/tests/test_operator_registry.py +++ b/tests/test_operator_registry.py @@ -176,11 +176,10 @@ def __init__(self, implementation_name): record_operator_binding("block-scaled FP8 Linear", linear) record_operator_binding("block-scaled FP8 Linear", linear_fallback) - log_operator_implementations(3) + log_operator_implementations() log_info.assert_called_once_with( - "Operator implementations (rank {}):\n{}", - 3, + "Operator implementations:\n{}", " Attention: triton\n" " block-scaled FP8 Linear: flashinfer_sm90, triton", ) diff --git a/tests/test_tp_rpc.py b/tests/test_tp_rpc.py index 15c27937..c5a05448 100644 --- a/tests/test_tp_rpc.py +++ b/tests/test_tp_rpc.py @@ -100,14 +100,17 @@ def test_prefix_offload_release_rpcs_use_failure_synchronized_world_path(): def test_operator_implementation_log_is_aligned_and_failure_synchronized(): - runner = object.__new__(ModelRunner) - runner.parallel_context = SimpleNamespace(world_rank=2) + 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(runner) + 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(2) + log_implementations.assert_called_once_with() def test_prefix_offload_release_rpc_surfaces_local_failure_after_status_sync(): From a984c9724d90dc6082d16383256b3e4c2d144074 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 19:14:57 +0800 Subject: [PATCH 09/35] perf: fuse qwen3.6 model projections --- src/sparsevllm/layers/linear.py | 4 +- src/sparsevllm/models/qwen3_5.py | 97 +++++++------- src/sparsevllm/models/qwen3_5_moe.py | 42 ++++-- src/sparsevllm/triton_kernel/moe_topk.py | 4 +- .../triton_kernel/qwen3_5/gated_shared_add.py | 80 +++++++++++ tests/test_qwen35_gated_shared_add.py | 21 +++ tests/test_qwen35_mixed_runtime.py | 124 +++++++++++++++++- 7 files changed, 308 insertions(+), 64 deletions(-) create mode 100644 src/sparsevllm/triton_kernel/qwen3_5/gated_shared_add.py create mode 100644 tests/test_qwen35_gated_shared_add.py 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/models/qwen3_5.py b/src/sparsevllm/models/qwen3_5.py index e7b82174..3397928b 100644 --- a/src/sparsevllm/models/qwen3_5.py +++ b/src/sparsevllm/models/qwen3_5.py @@ -388,39 +388,15 @@ def __init__(self, config) -> None: quantization = getattr(config, "quantization_config", None) hidden_size = int(config.hidden_size) - self.in_proj_q = ColumnParallelLinear( + self.in_proj_qkvz = MergedColumnParallelLinear( 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( - 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( + self.in_proj_ba = MergedColumnParallelLinear( hidden_size, - self.total_num_v_heads, - bias=False, - quantization=None, - ) - self.in_proj_a = ColumnParallelLinear( - hidden_size, - self.total_num_v_heads, + [self.total_num_v_heads, self.total_num_v_heads], bias=False, quantization=None, ) @@ -472,16 +448,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, @@ -495,9 +476,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( @@ -556,10 +543,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( @@ -579,8 +571,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: @@ -595,13 +587,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 @@ -823,7 +813,13 @@ def forward(self, positions: torch.Tensor, hidden_states: torch.Tensor) -> torch class Qwen35MLP(nn.Module): - def __init__(self, config, *, intermediate_size: int | None = None) -> 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 @@ -840,6 +836,7 @@ def __init__(self, config, *, intermediate_size: int | None = None) -> None: 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}.") @@ -951,6 +948,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 index db2a8cb6..9a6cd885 100644 --- a/src/sparsevllm/models/qwen3_5_moe.py +++ b/src/sparsevllm/models/qwen3_5_moe.py @@ -25,6 +25,7 @@ resolve_moe_router_provider, ) from sparsevllm.platforms import device_runtime +from sparsevllm.triton_kernel.qwen3_5.gated_shared_add import gated_shared_add from sparsevllm.utils.log import logger from sparsevllm.utils.weight_target import WeightTarget @@ -64,18 +65,37 @@ def __init__(self, config) -> None: ) self.provider = resolve_moe_router_provider(self.op_spec) self.weight = nn.Parameter( - torch.empty(self.num_experts, self.hidden_size) + 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]: - router_logits = F.linear(hidden_states, self.weight) + ) -> 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 + return topk_weights, topk_ids, shared_gate_logits class Qwen35MoePackedExperts(Qwen3MoePackedExperts): @@ -233,9 +253,7 @@ def __init__(self, config) -> None: self.shared_expert = Qwen35MLP( config, intermediate_size=int(config.shared_expert_intermediate_size), - ) - self.shared_expert_gate = nn.Linear( - int(config.hidden_size), 1, bias=False + reduce_results=False, ) def _forward_chunk( @@ -243,12 +261,12 @@ def _forward_chunk( hidden_states: torch.Tensor, ) -> torch.Tensor: shared_output = self.shared_expert(hidden_states) - topk_weights, topk_ids = self.gate(hidden_states) + topk_weights, topk_ids, shared_gate_logits = self.gate(hidden_states) local_output = self.experts(hidden_states, topk_ids, topk_weights) - routed_output = self.parallel_context.world_all_reduce(local_output) - shared_gate = torch.sigmoid(self.shared_expert_gate(hidden_states)) - gated_shared_output = shared_gate * shared_output - return routed_output + gated_shared_output + routed_output, shared_output = self.parallel_context.world_all_reduce( + torch.stack((local_output, shared_output)) + ) + return gated_shared_add(routed_output, shared_output, shared_gate_logits) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if hidden_states.dim() != 2: diff --git a/src/sparsevllm/triton_kernel/moe_topk.py b/src/sparsevllm/triton_kernel/moe_topk.py index ecfec960..e9d1f821 100644 --- a/src/sparsevllm/triton_kernel/moe_topk.py +++ b/src/sparsevllm/triton_kernel/moe_topk.py @@ -198,8 +198,8 @@ 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.") num_experts = int(router_logits.shape[1]) 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..0cc67525 --- /dev/null +++ b/src/sparsevllm/triton_kernel/qwen3_5/gated_shared_add.py @@ -0,0 +1,80 @@ +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 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.") + + 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/tests/test_qwen35_gated_shared_add.py b/tests/test_qwen35_gated_shared_add.py new file mode 100644 index 00000000..635ece37 --- /dev/null +++ b/tests/test_qwen35_gated_shared_add.py @@ -0,0 +1,21 @@ +import pytest +import torch + +from sparsevllm.triton_kernel.qwen3_5.gated_shared_add import gated_shared_add + + +pytestmark = 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) diff --git a/tests/test_qwen35_mixed_runtime.py b/tests/test_qwen35_mixed_runtime.py index 24ae8118..1b872458 100644 --- a/tests/test_qwen35_mixed_runtime.py +++ b/tests/test_qwen35_mixed_runtime.py @@ -1,7 +1,7 @@ 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 @@ -40,10 +40,12 @@ 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.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 +98,126 @@ def _make_config(tmp_path, **kwargs): return Config(model=str(tmp_path), **kwargs) +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_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(), + ) + + class _ResidentAdmissionCache: def __init__(self): self.num_free_slots = 1_000_000 From 4eb69e9c078604878fe901825a9fc928efd31143 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 19:15:17 +0800 Subject: [PATCH 10/35] perf: tune qwen3.6 fp8 moe kernels --- benchmark/tune_fp8_moe.py | 146 ++++++++++++ src/sparsevllm/operators/fp8_linear.py | 43 +--- src/sparsevllm/operators/moe.py | 97 ++++++++ src/sparsevllm/triton_kernel/moe.py | 180 ++++++++++---- src/sparsevllm/triton_kernel/moe_config.py | 79 ++++++ tests/test_moe_config.py | 63 +++++ tests/test_operator_providers.py | 264 ++++++++++++++++++--- tests/test_triton_moe.py | 29 +++ 8 files changed, 789 insertions(+), 112 deletions(-) create mode 100644 benchmark/tune_fp8_moe.py diff --git a/benchmark/tune_fp8_moe.py b/benchmark/tune_fp8_moe.py new file mode 100644 index 00000000..df0a74a7 --- /dev/null +++ b/benchmark/tune_fp8_moe.py @@ -0,0 +1,146 @@ +import argparse +import json +from pathlib import Path + +import torch +import triton + +from sparsevllm.triton_kernel.moe import ( + _prepare_expert_assignment, + _routed_fp8_gemm, +) +from sparsevllm.triton_kernel.moe_config import MoeGemmConfig + + +def parse_args(): + parser = argparse.ArgumentParser(description="Tune Qwen3.6 EP2 FP8 routed GEMMs.") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--local-assignments", type=int, default=4) + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--rep", type=int, default=300) + return parser.parse_args() + + +def candidates(): + for block_n, swap_ab, stages in ( + (64, True, range(2, 6)), + (64, False, range(2, 5)), + (128, True, range(2, 5)), + (128, False, range(2, 5)), + (32, True, range(3, 5)), + ): + yield from ( + MoeGemmConfig(16, block_n, 128, 1, 4, stage, swap_ab) + for stage in stages + ) + + +def run_stage(stage, inputs, weights, scales, topk_weights, alignment, args): + multiply_routing_weight = stage == "w2" + input_top_k = 1 if multiply_routing_weight else 8 + reference_config = MoeGemmConfig(16, 128, 128, 1, 4, 3) + reference = torch.empty(8, weights.shape[1], device="cuda", dtype=torch.bfloat16) + _routed_fp8_gemm( + inputs, + weights, + scales, + reference, + topk_weights, + alignment, + input_top_k=input_top_k, + multiply_routing_weight=multiply_routing_weight, + config=reference_config, + ) + torch.cuda.synchronize() + records = [] + for config in candidates(): + output = torch.empty_like(reference) + + def launch(): + _routed_fp8_gemm( + inputs, + weights, + scales, + output, + topk_weights, + alignment, + input_top_k=input_top_k, + multiply_routing_weight=multiply_routing_weight, + config=config, + ) + + record = { + "stage": stage, + "local_assignments": args.local_assignments, + **config.__dict__, + } + try: + launch() + torch.cuda.synchronize() + count = args.local_assignments + actual, expected = output[:count].float(), reference[:count].float() + record.update( + status="success", + max_abs_error=float((actual - expected).abs().max()), + latency_us=1000 + * triton.testing.do_bench( + launch, + warmup=args.warmup, + rep=args.rep, + return_mode="median", + ), + ) + except Exception as error: + record.update(status="invalid_config", error=f"{type(error).__name__}: {error}") + records.append(record) + print(json.dumps(record, sort_keys=True), flush=True) + return records + + +def main(): + args = parse_args() + if args.warmup <= 0 or args.rep <= 0: + raise ValueError("--warmup and --rep must be positive.") + if not 1 <= args.local_assignments <= 8: + raise ValueError("--local-assignments must be in [1, 8].") + torch.manual_seed(0) + hidden = torch.randn(1, 2048, device="cuda", dtype=torch.bfloat16) + activated = torch.randn(8, 512, device="cuda", dtype=torch.bfloat16) + expert_ids = list(range(args.local_assignments)) + list( + range(128, 136 - args.local_assignments) + ) + topk_ids = torch.tensor([expert_ids], device="cuda", dtype=torch.int32) + topk_weights = torch.full((1, 8), 0.125, device="cuda", dtype=torch.bfloat16) + alignment = _prepare_expert_assignment( + topk_ids, + block_size=16, + num_experts=256, + local_expert_start=0, + local_expert_end=128, + ) + shapes = {"w13": (hidden, 1024, 2048), "w2": (activated, 2048, 512)} + records = [] + for stage, (inputs, output_size, input_size) in shapes.items(): + weights = torch.randn( + 8, output_size, input_size, device="cuda", dtype=torch.bfloat16 + ).to(torch.float8_e4m3fn) + scales = torch.ones( + 8, output_size // 128, input_size // 128, device="cuda", dtype=torch.bfloat16 + ) + records.extend( + run_stage( + stage, + inputs, + weights, + scales, + topk_weights, + alignment, + args, + ) + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(records, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/src/sparsevllm/operators/fp8_linear.py b/src/sparsevllm/operators/fp8_linear.py index 33d43d7a..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,13 +51,6 @@ class FlashInferSm90Fp8LinearProvider(Fp8LinearProvider): name = "flashinfer_sm90" priority = 100 - def __init__(self) -> None: - self._fallback: TritonFp8LinearProvider | None = None - - @property - def implementation_name(self) -> str: - return self._fallback.name if self._fallback is not None else self.name - @classmethod def supports(cls, spec: Fp8LinearSpec, caps: DeviceCaps) -> SupportResult: if spec.block_shape != (128, 128): @@ -92,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}." @@ -101,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/moe.py b/src/sparsevllm/operators/moe.py index e0d212fd..0123950b 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -504,6 +504,103 @@ 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_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 != "NVIDIA H100 80GB HBM3": + return SupportResult.no( + "requires profiled NVIDIA H100 80GB HBM3 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, + ) + + def resolve_moe_provider( spec: MoeOpSpec, *, 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..60040cf0 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 { @@ -156,6 +157,52 @@ def _stage_table( } +_FP8_N64_SWAP = MoeGemmConfig(16, 64, 128, 1, 4, 3, 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) + + +# Qwen3.6-35B-A3B block-FP8 decode profiles tuned offline on H100. Larger +# token buckets retain the explicit generic configuration until profiled. +_TUNED_FP8_ROUTED_CONFIGS = { + MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.float8_e4m3fn, + 8, + 256, + 2048, + 512, + ): { + "w13": { + 1: _FP8_N64_SWAP, + 2: _FP8_N64_SWAP_S4, + 4: _FP8_N64_SWAP, + 8: _FP8_N128, + }, + "w2": { + 1: _FP8_N64_SWAP_S4, + 2: _FP8_N64_SWAP, + 4: _FP8_N128, + 8: _FP8_N64_SWAP, + }, + }, + MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.float8_e4m3fn, + 8, + 128, + 2048, + 512, + ): { + "w13": {1: _FP8_N64_SWAP_S5}, + "w2": {1: _FP8_N64_SWAP_S4}, + }, +} + + @lru_cache(maxsize=None) def _resolve_moe_gemm_config( dtype: torch.dtype, @@ -234,3 +281,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/tests/test_moe_config.py b/tests/test_moe_config.py index 87e38c72..f7deafbf 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, ) @@ -155,3 +156,65 @@ 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( + ("stage", "tokens", "block_n", "num_stages", "swap_ab"), + [ + ("w13", 1, 64, 3, True), + ("w13", 2, 64, 4, True), + ("w13", 8, 128, 3, False), + ("w2", 1, 64, 4, True), + ("w2", 4, 128, 3, False), + ("w2", 8, 64, 3, True), + ], +) +def test_h100_qwen36_fp8_routed_config(stage, tokens, block_n, num_stages, swap_ab): + config = resolve_fp8_routed_gemm_config( + num_tokens=tokens, + top_k=8, + num_local_experts=256, + hidden_size=2048, + intermediate_size=512, + stage=stage, + device_name="NVIDIA H100 80GB HBM3", + device_capability=(9, 0), + ) + + assert (config.block_n, config.num_stages, config.swap_ab) == ( + block_n, + num_stages, + swap_ab, + ) + + +def test_h100_qwen36_fp8_ep2_uses_profiled_configs(): + common = dict( + num_tokens=1, + top_k=8, + num_local_experts=128, + hidden_size=2048, + intermediate_size=512, + device_name="NVIDIA H100 80GB HBM3", + device_capability=(9, 0), + ) + w13 = resolve_fp8_routed_gemm_config(**common, stage="w13") + w2 = resolve_fp8_routed_gemm_config(**common, stage="w2") + + assert (w13.block_n, w13.num_stages, w13.swap_ab) == (64, 5, True) + assert (w2.block_n, w2.num_stages, w2.swap_ab) == (64, 4, True) + + +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 d53479f6..16687244 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -12,7 +12,13 @@ TritonFp8LinearProvider, resolve_fp8_linear_provider, ) -from sparsevllm.operators.moe import MOE_REGISTRY, MoeOpSpec, resolve_moe_provider +from sparsevllm.operators.moe import ( + MOE_REGISTRY, + FlashInferCutlassFp8MoeProvider, + HopperQwen36HybridFp8MoeProvider, + MoeOpSpec, + resolve_moe_provider, +) from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum @@ -64,6 +70,7 @@ def _moe_spec( tp_size=1, routing_method="softmax", scale_dtype=None, + cuda_graph=True, ) -> MoeOpSpec: return MoeOpSpec( num_experts=num_experts, @@ -75,7 +82,7 @@ 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, @@ -155,9 +162,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 +186,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,15 +198,13 @@ 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() - assert provider.implementation_name == "flashinfer_sm90" x = torch.ones(2, 128, dtype=torch.bfloat16) weight = torch.ones(128, 128).to(torch.float8_e4m3fn) scale = torch.ones(1, 1) @@ -213,21 +218,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 - assert provider.implementation_name == "triton" def test_flashinfer_linear_does_not_mask_other_runtime_failures(): @@ -246,14 +241,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( @@ -522,6 +513,225 @@ def test_fp8_moe_prefers_flashinfer_only_on_sm90(): assert blackwell.provider.name == "triton" +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" + + +@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_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_triton_moe.py b/tests/test_triton_moe.py index acd90195..1ab64422 100644 --- a/tests/test_triton_moe.py +++ b/tests/test_triton_moe.py @@ -5,6 +5,7 @@ import torch.nn.functional as F from sparsevllm.triton_kernel.moe import ( + _prepare_expert_assignment, fused_moe, fused_moe_gate_up_swiglu, moe_align_block_size, @@ -71,6 +72,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 +155,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") From 5de1f0567bf5a173ba2307321fe0da4f68c5910f Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 19:15:28 +0800 Subject: [PATCH 11/35] perf: accelerate hopper tp2 all-reduce --- .../distributed/parallel_context.py | 31 ++++++- src/sparsevllm/operators/all_reduce.py | 84 +++++++++++++++++++ tests/test_all_reduce_operator.py | 17 ++++ tests/test_parallel_context.py | 28 ++++++- 4 files changed, 153 insertions(+), 7 deletions(-) create mode 100644 src/sparsevllm/operators/all_reduce.py create mode 100644 tests/test_all_reduce_operator.py diff --git a/src/sparsevllm/distributed/parallel_context.py b/src/sparsevllm/distributed/parallel_context.py index 1304c8d5..d418e4bf 100644 --- a/src/sparsevllm/distributed/parallel_context.py +++ b/src/sparsevllm/distributed/parallel_context.py @@ -1,10 +1,12 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field, replace import torch import torch.distributed as dist +from sparsevllm.operators.all_reduce import AllReduceProvider, resolve_all_reduce_provider + 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)) @@ -155,6 +157,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 +237,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( @@ -390,7 +396,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 +410,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/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/tests/test_all_reduce_operator.py b/tests/test_all_reduce_operator.py new file mode 100644 index 00000000..b46fb5e8 --- /dev/null +++ b/tests/test_all_reduce_operator.py @@ -0,0 +1,17 @@ +from unittest.mock import Mock + +import torch + +from sparsevllm.operators.all_reduce import HopperTp2FlashInferAllReduceProvider + + +def test_flashinfer_all_reduce_dispatches_unsupported_shape_before_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) diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py index 76bfa1bc..a914e7b5 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch @@ -124,9 +124,10 @@ 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, "new_group", side_effect=lambda _ranks: object()), + 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, @@ -156,6 +157,7 @@ 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) @@ -371,6 +373,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( From 2fabb96e9193ab0d5289f49d29e2b8708778eba3 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 19:15:38 +0800 Subject: [PATCH 12/35] test: add reproducible vllm benchmark --- benchmark/microbench.py | 10 +- benchmark/vllm_microbench.py | 282 +++++++++++++++++++++++++++++ tests/test_microbench_artifacts.py | 8 +- tests/test_vllm_microbench.py | 32 ++++ 4 files changed, 327 insertions(+), 5 deletions(-) create mode 100644 benchmark/vllm_microbench.py create mode 100644 tests/test_vllm_microbench.py diff --git a/benchmark/microbench.py b/benchmark/microbench.py index ffa5bf76..a1e0b3b0 100644 --- a/benchmark/microbench.py +++ b/benchmark/microbench.py @@ -286,6 +286,8 @@ def _artifact_records(args, rows: list[dict[str, Any]]) -> list[dict[str, Any]]: record.setdefault("itl_ms", row["itl"]) if "mem" in row: record.setdefault("peak_memory_gb", row["mem"]) + if "duration_s" in row: + record.setdefault("e2e_latency_s", row["duration_s"]) records.append(record) return records @@ -345,18 +347,19 @@ def _write_output_dir(args, rows: list[dict[str, Any]]) -> None: f"- Batch sizes: `{args.batch_sizes}`", f"- Output length: `{args.output_len}`", "", - "| Method | Prompt tokens | Batch | Status | TTFT s | Prefill tok/s | Decode tok/s | Peak GB | Decode speedup |", - "| --- | ---: | ---: | --- | ---: | ---: | ---: | ---: | ---: |", + "| Method | Prompt tokens | Batch | Status | E2E s | TTFT s | Prefill tok/s | Decode tok/s | Peak GB | Decode speedup |", + "| --- | ---: | ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: |", ] for record in records: ok = record["status"] == "success" speedup = record.get("speedup_vs_vanilla_decode") report_lines.append( - "| {method} | {prompt} | {batch} | {status} | {ttft} | {prefill} | {decode} | {mem} | {speedup} |".format( + "| {method} | {prompt} | {batch} | {status} | {e2e} | {ttft} | {prefill} | {decode} | {mem} | {speedup} |".format( method=record.get("method", ""), prompt=record.get("prompt_tokens", ""), batch=record.get("batch_size", ""), status=record["status"], + e2e=f"{record.get('e2e_latency_s', 0.0):.3f}" if ok else "", ttft=f"{record.get('ttft_s', 0.0):.3f}" if ok else "", prefill=f"{record.get('prefill_tok_s', 0.0):.1f}" if ok else "", decode=f"{record.get('decode_tok_s', 0.0):.1f}" if ok else "", @@ -754,6 +757,7 @@ def add_wave(max_new_requests: int): "prefill_tp": prefill_tp, "decode_tp": decode_tp, "ttft": ttft, + "duration_s": duration, "itl": avg_itl, "avg_bs": avg_active_bs, "mem": peak_mem, 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/tests/test_microbench_artifacts.py b/tests/test_microbench_artifacts.py index 3bdc0c99..f4cbdf27 100644 --- a/tests/test_microbench_artifacts.py +++ b/tests/test_microbench_artifacts.py @@ -165,7 +165,7 @@ def test_benchmark_sparse_method_preserves_graph_enabling_legacy_alias(method): @pytest.mark.parametrize("enabled", [False, True]) -def test_artifact_records_include_step_timing_mode(enabled): +def test_artifact_records_include_step_timing_mode_and_e2e_latency(enabled): args = SimpleNamespace( output_len=8, temperature=0.0, @@ -173,9 +173,13 @@ def test_artifact_records_include_step_timing_mode(enabled): synchronize_step_timing=enabled, ) - records = _artifact_records(args, [{"status": "SUCCESS", "length": 16}]) + records = _artifact_records( + args, + [{"status": "SUCCESS", "length": 16, "duration_s": 1.25}], + ) assert records[0]["synchronize_step_timing"] is enabled + assert records[0]["e2e_latency_s"] == 1.25 def test_output_metadata_records_step_timing_mode(tmp_path, monkeypatch): diff --git a/tests/test_vllm_microbench.py b/tests/test_vllm_microbench.py new file mode 100644 index 00000000..41d20985 --- /dev/null +++ b/tests/test_vllm_microbench.py @@ -0,0 +1,32 @@ +import argparse +from types import SimpleNamespace + +import pytest + +from benchmark.vllm_microbench import _parse_positive_ints, _validate_args + + +def test_vllm_microbench_parses_unique_batch_sizes(): + assert _parse_positive_ints("1,2,4") == [1, 2, 4] + + +@pytest.mark.parametrize("value", ["", "0,1", "1,1"]) +def test_vllm_microbench_rejects_invalid_batch_sizes(value): + with pytest.raises(argparse.ArgumentTypeError): + _parse_positive_ints(value) + + +def test_vllm_microbench_rejects_short_model_context(): + args = SimpleNamespace( + input_len=1024, + output_len=128, + num_warmups=2, + num_iters=5, + tensor_parallel_size=2, + max_model_len=1151, + max_num_batched_tokens=4096, + gpu_memory_utilization=0.7, + ) + + with pytest.raises(ValueError, match=r"input_len \+ output_len"): + _validate_args(args) From bf6097a49b23419153e62bedff89c4f3fc8d6f49 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 19:51:10 +0800 Subject: [PATCH 13/35] perf: tune qwen3.6 fp8 moe for h20 --- src/sparsevllm/operators/moe.py | 13 ++++++++-- src/sparsevllm/triton_kernel/moe_config.py | 25 ++++++++++++++++++ tests/test_moe_config.py | 30 ++++++++++++++++++++++ tests/test_operator_providers.py | 24 +++++++++++++++++ 4 files changed, 90 insertions(+), 2 deletions(-) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 0123950b..50a733eb 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -510,6 +510,7 @@ class HopperQwen36HybridFp8MoeProvider(FlashInferCutlassFp8MoeProvider): name = "hopper_qwen36_hybrid_fp8" priority = 110 + PROFILED_DEVICE_NAME = "NVIDIA H100 80GB HBM3" PROFILED_SHAPES = frozenset( { (256, 256, 2048, 512, 8, 1, 1), @@ -522,9 +523,9 @@ class HopperQwen36HybridFp8MoeProvider(FlashInferCutlassFp8MoeProvider): 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 != "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}" ) actual_shape = ( @@ -601,6 +602,14 @@ def run( ) +@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: 1, 2: 1} + + def resolve_moe_provider( spec: MoeOpSpec, *, diff --git a/src/sparsevllm/triton_kernel/moe_config.py b/src/sparsevllm/triton_kernel/moe_config.py index 60040cf0..39e0759c 100644 --- a/src/sparsevllm/triton_kernel/moe_config.py +++ b/src/sparsevllm/triton_kernel/moe_config.py @@ -161,11 +161,36 @@ def _stage_table( _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_S4 = MoeGemmConfig(16, 128, 128, 1, 4, 4, True) # Qwen3.6-35B-A3B block-FP8 decode profiles tuned offline on H100. Larger # token buckets retain the explicit generic configuration until profiled. _TUNED_FP8_ROUTED_CONFIGS = { + MoeGemmShape( + "H20", + (9, 0), + torch.float8_e4m3fn, + 8, + 256, + 2048, + 512, + ): { + "w13": {1: _FP8_N64_SWAP_S4}, + "w2": {1: _FP8_N128_SWAP_S4}, + }, + MoeGemmShape( + "H20", + (9, 0), + torch.float8_e4m3fn, + 8, + 128, + 2048, + 512, + ): { + "w13": {1: _FP8_N64_SWAP_S4}, + "w2": {1: _FP8_N64_SWAP}, + }, MoeGemmShape( "NVIDIA H100 80GB HBM3", (9, 0), diff --git a/tests/test_moe_config.py b/tests/test_moe_config.py index f7deafbf..fc76d975 100644 --- a/tests/test_moe_config.py +++ b/tests/test_moe_config.py @@ -205,6 +205,36 @@ def test_h100_qwen36_fp8_ep2_uses_profiled_configs(): assert (w2.block_n, w2.num_stages, w2.swap_ab) == (64, 4, True) +@pytest.mark.parametrize( + ("local_experts", "stage", "block_n", "num_stages"), + [ + (256, "w13", 64, 4), + (256, "w2", 128, 4), + (128, "w13", 64, 4), + (128, "w2", 64, 3), + ], +) +def test_h20_qwen36_fp8_uses_profiled_decode_configs( + local_experts, stage, block_n, num_stages +): + config = resolve_fp8_routed_gemm_config( + num_tokens=1, + top_k=8, + num_local_experts=local_experts, + hidden_size=2048, + intermediate_size=512, + stage=stage, + device_name="NVIDIA H20", + device_capability=(9, 0), + ) + + assert (config.block_n, config.num_stages, config.swap_ab) == ( + block_n, + num_stages, + True, + ) + + def test_fp8_routed_unknown_shape_uses_explicit_default(): config = resolve_fp8_routed_gemm_config( num_tokens=1, diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 16687244..4c724974 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -15,6 +15,7 @@ from sparsevllm.operators.moe import ( MOE_REGISTRY, FlashInferCutlassFp8MoeProvider, + H20Qwen36HybridFp8MoeProvider, HopperQwen36HybridFp8MoeProvider, MoeOpSpec, resolve_moe_provider, @@ -610,6 +611,29 @@ def test_qwen36_hybrid_moe_supports_eager_execution(): 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_h20_qwen36_hybrid_moe_limits_triton_to_profiled_token_count(): + assert H20Qwen36HybridFp8MoeProvider.TRITON_MAX_TOKENS_BY_EP_SIZE == {1: 1, 2: 1} + + def test_qwen36_hybrid_moe_dispatches_by_token_bucket(): provider = HopperQwen36HybridFp8MoeProvider() spec = _moe_spec( From d6b11fdad2d36e1a471db8934c7eea09213347e7 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 20:12:46 +0800 Subject: [PATCH 14/35] test: add bf16 moe tuner --- benchmark/tune_bf16_moe.py | 183 +++++++++++++++++++++++++++++++++++++ 1 file changed, 183 insertions(+) create mode 100644 benchmark/tune_bf16_moe.py diff --git a/benchmark/tune_bf16_moe.py b/benchmark/tune_bf16_moe.py new file mode 100644 index 00000000..3d518a52 --- /dev/null +++ b/benchmark/tune_bf16_moe.py @@ -0,0 +1,183 @@ +import argparse +import json +from pathlib import Path + +import torch +import triton + +from sparsevllm.triton_kernel.moe import ( + _prepare_expert_assignment, + _routed_gate_up_swiglu, + _routed_gemm, +) +from sparsevllm.triton_kernel.moe_config import MoeGemmConfig +from sparsevllm.triton_kernel.silu_and_mul import silu_and_mul_fwd + + +def parse_args(): + parser = argparse.ArgumentParser(description="Tune Qwen3.6 BF16 routed GEMMs.") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--rep", type=int, default=300) + return parser.parse_args() + + +def candidates(): + for block_n in (32, 64, 128): + for block_k, warps, stages in ( + (32, 4, 3), + (32, 4, 4), + (64, 4, 2), + (64, 4, 3), + (64, 4, 4), + (64, 8, 3), + ): + yield MoeGemmConfig(16, block_n, block_k, 8, warps, stages) + + +def benchmark(stage, launch, output, reference, config, args): + record = {"stage": stage, **config.__dict__} + try: + launch() + torch.cuda.synchronize() + actual, expected = output.float(), reference.float() + record.update( + status="success", + max_abs_error=float((actual - expected).abs().max()), + relative_l2_error=float( + torch.linalg.vector_norm(actual - expected) + / torch.linalg.vector_norm(expected) + ), + latency_us=1000 + * triton.testing.do_bench( + launch, + warmup=args.warmup, + rep=args.rep, + return_mode="median", + ), + ) + except Exception as error: + record.update(status="invalid_config", error=f"{type(error).__name__}: {error}") + print(json.dumps(record, sort_keys=True), flush=True) + return record + + +def main(): + args = parse_args() + if args.warmup <= 0 or args.rep <= 0: + raise ValueError("--warmup and --rep must be positive.") + torch.manual_seed(0) + hidden = torch.randn(1, 2048, device="cuda", dtype=torch.bfloat16) + w13 = torch.randn(8, 1024, 2048, device="cuda", dtype=torch.bfloat16) + w2 = torch.randn(8, 2048, 512, device="cuda", dtype=torch.bfloat16) + topk_ids = torch.arange(8, device="cuda", dtype=torch.int32).view(1, 8) + topk_weights = torch.full((1, 8), 0.125, device="cuda", dtype=torch.bfloat16) + alignment = _prepare_expert_assignment( + topk_ids, + block_size=16, + num_experts=256, + local_expert_start=0, + local_expert_end=256, + ) + reference_config = MoeGemmConfig(16, 128, 32, 8, 4, 4).as_triton_kwargs() + w13_reference = torch.empty(8, 1024, device="cuda", dtype=torch.bfloat16) + _routed_gemm( + hidden, + w13, + w13_reference, + topk_weights, + alignment, + input_top_k=8, + multiply_routing_weight=False, + launch_config=reference_config, + ) + activated = silu_and_mul_fwd(w13_reference.clone()) + w2_reference = torch.empty(8, 2048, device="cuda", dtype=torch.bfloat16) + _routed_gemm( + activated, + w2, + w2_reference, + topk_weights, + alignment, + input_top_k=1, + multiply_routing_weight=True, + launch_config=reference_config, + ) + torch.cuda.synchronize() + + records = [] + for config in candidates(): + launch_config = config.as_triton_kwargs() + w13_output = torch.empty_like(w13_reference) + records.append( + benchmark( + "w13", + lambda: _routed_gemm( + hidden, + w13, + w13_output, + topk_weights, + alignment, + input_top_k=8, + multiply_routing_weight=False, + launch_config=launch_config, + ), + w13_output, + w13_reference, + config, + args, + ) + ) + fused_output = torch.empty_like(activated) + records.append( + benchmark( + "gate_up_swiglu", + lambda: _routed_gate_up_swiglu( + hidden, + w13, + fused_output, + alignment, + input_top_k=8, + launch_config=launch_config, + ), + fused_output, + activated, + config, + args, + ) + ) + w2_output = torch.empty_like(w2_reference) + records.append( + benchmark( + "w2", + lambda: _routed_gemm( + activated, + w2, + w2_output, + topk_weights, + alignment, + input_top_k=1, + multiply_routing_weight=True, + launch_config=launch_config, + ), + w2_output, + w2_reference, + config, + args, + ) + ) + + result = { + "device": torch.cuda.get_device_name(), + "torch_version": torch.__version__, + "triton_version": triton.__version__, + "warmup": args.warmup, + "rep": args.rep, + "records": records, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n") + + +if __name__ == "__main__": + main() From 2b0643fe45bb801556d5e7f7f8791e87e52e0c83 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 20:22:47 +0800 Subject: [PATCH 15/35] perf: tune qwen3.6 bf16 moe for h20 --- src/sparsevllm/operators/moe.py | 21 +++++++++++----- src/sparsevllm/triton_kernel/moe_config.py | 29 +++++++++++++++++----- tests/test_moe_config.py | 26 +++++++++++++++++++ tests/test_operator_providers.py | 29 ++++++++++++++++++++++ 4 files changed, 93 insertions(+), 12 deletions(-) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 50a733eb..5cdffdec 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -356,6 +356,8 @@ class TritonHopperFusedMoeProvider(MoeProvider): name = "triton_hopper_fused" priority = 20 gate_up_order = "gate_up" + PROFILED_DEVICE_NAME = "NVIDIA H100 80GB HBM3" + PROFILED_SHAPE = (128, 64, 2048, 384, 8, 2, 2) @classmethod def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: @@ -373,14 +375,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 +391,10 @@ def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: spec.tp_size, spec.ep_size, ) - if actual_shape != profiled_shape: + if actual_shape != cls.PROFILED_SHAPE: return SupportResult.no( - "requires profiled TP2xEP2 MoE shape " - f"{profiled_shape}, got {actual_shape}" + "requires profiled MoE shape " + f"{cls.PROFILED_SHAPE}, got {actual_shape}" ) return SupportResult.yes() @@ -427,6 +428,14 @@ def run( ) +@MOE_REGISTRY.register +class H20Qwen36FusedMoeProvider(TritonHopperFusedMoeProvider): + name = "h20_qwen36_fused_bf16" + priority = 21 + PROFILED_DEVICE_NAME = "NVIDIA H20" + PROFILED_SHAPE = (256, 256, 2048, 512, 8, 1, 1) + + @MOE_REGISTRY.register class TritonMoeProvider(MoeProvider): name = "triton" diff --git a/src/sparsevllm/triton_kernel/moe_config.py b/src/sparsevllm/triton_kernel/moe_config.py index 39e0759c..32cad6e3 100644 --- a/src/sparsevllm/triton_kernel/moe_config.py +++ b/src/sparsevllm/triton_kernel/moe_config.py @@ -107,9 +107,13 @@ 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}, + "w2": {1: _H}, + }, 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), @@ -137,9 +141,18 @@ def _stage_table( } -# 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}, MoeGemmShape( "NVIDIA H100 80GB HBM3", (9, 0), @@ -253,7 +266,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, @@ -264,7 +279,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( diff --git a/tests/test_moe_config.py b/tests/test_moe_config.py index fc76d975..bea40b3c 100644 --- a/tests/test_moe_config.py +++ b/tests/test_moe_config.py @@ -106,6 +106,32 @@ def test_h20_qwen3_moe_config_is_shape_and_stage_aware(): assert large.block_m == 64 +def test_h20_qwen36_decode_uses_profiled_bf16_configs(): + common = dict( + dtype=torch.bfloat16, + num_tokens=1, + top_k=8, + num_local_experts=256, + hidden_size=2048, + intermediate_size=512, + device_name="NVIDIA H20", + device_capability=(9, 0), + ) + w13 = resolve_moe_gemm_config(**common, stage="w13") + fused = resolve_moe_gemm_config(**common, stage="gate_up_swiglu") + w2 = resolve_moe_gemm_config(**common, stage="w2") + + assert (w13.block_n, w13.block_k, w13.num_stages) == (32, 64, 4) + assert fused == w13 + assert (w2.block_n, w2.block_k, w2.num_stages) == (32, 64, 3) + + unprofiled = resolve_moe_gemm_config( + **{**common, "num_tokens": 2}, + stage="w13", + ) + assert (unprofiled.block_n, unprofiled.block_k) == (128, 32) + + def test_fallback_heuristic_uses_logical_assignment_count(): common = dict( dtype=torch.float16, diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 4c724974..26270dd2 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -366,6 +366,35 @@ def test_hopper_fused_moe_uses_profiled_tp_ep_shape(): assert resolved.provider.name == "triton_hopper_fused" +def test_h20_qwen36_bf16_moe_uses_profiled_provider(): + spec = _moe_spec( + activation_dtype=torch.bfloat16, + weight_dtype=torch.bfloat16, + block_shape=None, + hidden_size=2048, + intermediate_size=512, + num_local_experts=256, + num_experts=256, + top_k=8, + ep_size=1, + tp_size=1, + ) + + 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" + + @pytest.mark.parametrize( ("tp_size", "ep_size", "intermediate_size", "num_local_experts"), [(4, 1, 384, 256), (2, 2, 768, 128), (1, 4, 1536, 64)], From d02d78ebacd37fea234bb0fbc39c6bbc03350a94 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 21:04:43 +0800 Subject: [PATCH 16/35] test: extend bf16 moe tuner --- benchmark/tune_bf16_moe.py | 47 +++++++++++++++++++++++++++++++++----- 1 file changed, 41 insertions(+), 6 deletions(-) diff --git a/benchmark/tune_bf16_moe.py b/benchmark/tune_bf16_moe.py index 3d518a52..04ffeefc 100644 --- a/benchmark/tune_bf16_moe.py +++ b/benchmark/tune_bf16_moe.py @@ -17,6 +17,9 @@ def parse_args(): parser = argparse.ArgumentParser(description="Tune Qwen3.6 BF16 routed GEMMs.") parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--intermediate-size", type=int, default=512) + parser.add_argument("--num-local-experts", type=int, default=256) + parser.add_argument("--local-assignments", type=int, default=8) parser.add_argument("--warmup", type=int, default=50) parser.add_argument("--rep", type=int, default=300) return parser.parse_args() @@ -40,7 +43,8 @@ def benchmark(stage, launch, output, reference, config, args): try: launch() torch.cuda.synchronize() - actual, expected = output.float(), reference.float() + actual = output[: args.local_assignments].float() + expected = reference[: args.local_assignments].float() record.update( status="success", max_abs_error=float((actual - expected).abs().max()), @@ -66,21 +70,49 @@ def main(): args = parse_args() if args.warmup <= 0 or args.rep <= 0: raise ValueError("--warmup and --rep must be positive.") + if args.intermediate_size <= 0 or not 1 <= args.num_local_experts <= 256: + raise ValueError("--intermediate-size and --num-local-experts must be valid.") + if not 1 <= args.local_assignments <= min(8, args.num_local_experts): + raise ValueError("--local-assignments must fit the local expert range.") + if args.num_local_experts + 8 - args.local_assignments > 256: + raise ValueError("Remote assignments must fit the global expert range.") torch.manual_seed(0) + intermediate_size = args.intermediate_size + num_weight_experts = args.local_assignments hidden = torch.randn(1, 2048, device="cuda", dtype=torch.bfloat16) - w13 = torch.randn(8, 1024, 2048, device="cuda", dtype=torch.bfloat16) - w2 = torch.randn(8, 2048, 512, device="cuda", dtype=torch.bfloat16) - topk_ids = torch.arange(8, device="cuda", dtype=torch.int32).view(1, 8) + w13 = torch.randn( + num_weight_experts, + 2 * intermediate_size, + 2048, + device="cuda", + dtype=torch.bfloat16, + ) + w2 = torch.randn( + num_weight_experts, + 2048, + intermediate_size, + device="cuda", + dtype=torch.bfloat16, + ) + expert_ids = list(range(args.local_assignments)) + list( + range( + args.num_local_experts, + args.num_local_experts + 8 - args.local_assignments, + ) + ) + topk_ids = torch.tensor([expert_ids], device="cuda", dtype=torch.int32) topk_weights = torch.full((1, 8), 0.125, device="cuda", dtype=torch.bfloat16) alignment = _prepare_expert_assignment( topk_ids, block_size=16, num_experts=256, local_expert_start=0, - local_expert_end=256, + local_expert_end=args.num_local_experts, ) reference_config = MoeGemmConfig(16, 128, 32, 8, 4, 4).as_triton_kwargs() - w13_reference = torch.empty(8, 1024, device="cuda", dtype=torch.bfloat16) + w13_reference = torch.empty( + 8, 2 * intermediate_size, device="cuda", dtype=torch.bfloat16 + ) _routed_gemm( hidden, w13, @@ -171,6 +203,9 @@ def main(): "device": torch.cuda.get_device_name(), "torch_version": torch.__version__, "triton_version": triton.__version__, + "intermediate_size": intermediate_size, + "num_local_experts": args.num_local_experts, + "local_assignments": args.local_assignments, "warmup": args.warmup, "rep": args.rep, "records": records, From bee4d203bb4b164b163f37da36f96b1517dbf473 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 21:08:46 +0800 Subject: [PATCH 17/35] perf: tune parallel bf16 moe for h20 --- src/sparsevllm/operators/moe.py | 14 +++++++---- src/sparsevllm/triton_kernel/moe_config.py | 27 ++++++++++++++++++++++ tests/test_moe_config.py | 26 +++++++++++++++++++++ tests/test_operator_providers.py | 19 +++++++++++---- 4 files changed, 76 insertions(+), 10 deletions(-) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 5cdffdec..e1a7c12d 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -357,7 +357,7 @@ class TritonHopperFusedMoeProvider(MoeProvider): priority = 20 gate_up_order = "gate_up" PROFILED_DEVICE_NAME = "NVIDIA H100 80GB HBM3" - PROFILED_SHAPE = (128, 64, 2048, 384, 8, 2, 2) + PROFILED_SHAPES = ((128, 64, 2048, 384, 8, 2, 2),) @classmethod def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: @@ -391,10 +391,10 @@ def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: spec.tp_size, spec.ep_size, ) - if actual_shape != cls.PROFILED_SHAPE: + if actual_shape not in cls.PROFILED_SHAPES: return SupportResult.no( - "requires profiled MoE shape " - f"{cls.PROFILED_SHAPE}, got {actual_shape}" + "requires a profiled MoE shape in " + f"{cls.PROFILED_SHAPES}, got {actual_shape}" ) return SupportResult.yes() @@ -433,7 +433,11 @@ class H20Qwen36FusedMoeProvider(TritonHopperFusedMoeProvider): name = "h20_qwen36_fused_bf16" priority = 21 PROFILED_DEVICE_NAME = "NVIDIA H20" - PROFILED_SHAPE = (256, 256, 2048, 512, 8, 1, 1) + 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 diff --git a/src/sparsevllm/triton_kernel/moe_config.py b/src/sparsevllm/triton_kernel/moe_config.py index 32cad6e3..521b6ff6 100644 --- a/src/sparsevllm/triton_kernel/moe_config.py +++ b/src/sparsevllm/triton_kernel/moe_config.py @@ -93,6 +93,7 @@ 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) def _stage_table( @@ -114,6 +115,14 @@ def _stage_table( "w13": {1: _G}, "w2": {1: _H}, }, + MoeGemmShape("H20", (9, 0), torch.bfloat16, 8, 256, 2048, 256): { + "w13": {1: _G}, + "w2": {1: _H}, + }, + MoeGemmShape("H20", (9, 0), torch.bfloat16, 8, 128, 2048, 512): { + "w13": {1: _G}, + "w2": {1: _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), @@ -153,6 +162,24 @@ def _stage_table( 2048, 512, ): {1: _G}, + MoeGemmShape( + "H20", + (9, 0), + torch.bfloat16, + 8, + 256, + 2048, + 256, + ): {1: _G}, + MoeGemmShape( + "H20", + (9, 0), + torch.bfloat16, + 8, + 128, + 2048, + 512, + ): {1: _G}, MoeGemmShape( "NVIDIA H100 80GB HBM3", (9, 0), diff --git a/tests/test_moe_config.py b/tests/test_moe_config.py index bea40b3c..722e9938 100644 --- a/tests/test_moe_config.py +++ b/tests/test_moe_config.py @@ -132,6 +132,32 @@ def test_h20_qwen36_decode_uses_profiled_bf16_configs(): assert (unprofiled.block_n, unprofiled.block_k) == (128, 32) +@pytest.mark.parametrize( + ("num_local_experts", "intermediate_size", "expected_w2"), + [(256, 256, (32, 64, 3)), (128, 512, (64, 64, 4))], +) +def test_h20_qwen36_parallel_decode_uses_profiled_bf16_configs( + num_local_experts, + intermediate_size, + expected_w2, +): + common = dict( + dtype=torch.bfloat16, + num_tokens=1, + top_k=8, + num_local_experts=num_local_experts, + hidden_size=2048, + intermediate_size=intermediate_size, + device_name="NVIDIA H20", + device_capability=(9, 0), + ) + + fused = resolve_moe_gemm_config(**common, stage="gate_up_swiglu") + w2 = resolve_moe_gemm_config(**common, stage="w2") + assert (fused.block_n, fused.block_k, fused.num_stages) == (32, 64, 4) + assert (w2.block_n, w2.block_k, w2.num_stages) == expected_w2 + + def test_fallback_heuristic_uses_logical_assignment_count(): common = dict( dtype=torch.float16, diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 26270dd2..47c043b5 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -366,18 +366,27 @@ def test_hopper_fused_moe_uses_profiled_tp_ep_shape(): assert resolved.provider.name == "triton_hopper_fused" -def test_h20_qwen36_bf16_moe_uses_profiled_provider(): +@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_h20_qwen36_bf16_moe_uses_profiled_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=512, - num_local_experts=256, + intermediate_size=intermediate_size, + num_local_experts=num_local_experts, num_experts=256, top_k=8, - ep_size=1, - tp_size=1, + ep_size=ep_size, + tp_size=tp_size, ) resolved = OpResolver(MOE_REGISTRY).resolve( From 347a3490d0c3243ddf6d2b04b47451d39f72ac57 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 21:32:25 +0800 Subject: [PATCH 18/35] perf: fuse h20 shared expert swiglu --- src/sparsevllm/models/qwen3_5.py | 30 ++- src/sparsevllm/operators/gate_up_swiglu.py | 156 ++++++++++++++++ .../triton_kernel/gate_up_swiglu.py | 176 ++++++++++++++++++ tests/test_gate_up_swiglu_operator.py | 85 +++++++++ 4 files changed, 444 insertions(+), 3 deletions(-) create mode 100644 src/sparsevllm/operators/gate_up_swiglu.py create mode 100644 src/sparsevllm/triton_kernel/gate_up_swiglu.py create mode 100644 tests/test_gate_up_swiglu_operator.py diff --git a/src/sparsevllm/models/qwen3_5.py b/src/sparsevllm/models/qwen3_5.py index 3397928b..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 @@ -840,11 +843,32 @@ def __init__( ) 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: diff --git a/src/sparsevllm/operators/gate_up_swiglu.py b/src/sparsevllm/operators/gate_up_swiglu.py new file mode 100644 index 00000000..b9deb222 --- /dev/null +++ b/src/sparsevllm/operators/gate_up_swiglu.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +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 + + +class GateUpProjection(Protocol): + weight: torch.Tensor + + def __call__(self, inputs: torch.Tensor) -> torch.Tensor: ... + + +@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 or not 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: GateUpProjection, + ) -> torch.Tensor: + raise NotImplementedError + + +GATE_UP_SWIGLU_REGISTRY: OpRegistry[ + GateUpSwiGLUOpSpec, GateUpSwiGLUProvider +] = OpRegistry("gate/up SwiGLU") + + +@GATE_UP_SWIGLU_REGISTRY.register +class TorchGateUpSwiGLUProvider(GateUpSwiGLUProvider): + name = "torch" + 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: GateUpProjection, + ) -> 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(TorchGateUpSwiGLUProvider): + 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: GateUpProjection, + ) -> torch.Tensor: + if inputs.shape[0] != 1: + return super().run(spec, inputs, projection) + from sparsevllm.triton_kernel.gate_up_swiglu import ( + gate_up_swiglu, + resolve_h20_gate_up_swiglu_config, + ) + + local_intermediate_size = spec.intermediate_size // spec.tp_size + return gate_up_swiglu( + inputs, + projection.weight, + resolve_h20_gate_up_swiglu_config( + inputs.shape[0], spec.hidden_size, local_intermediate_size + ), + ) + + +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/triton_kernel/gate_up_swiglu.py b/src/sparsevllm/triton_kernel/gate_up_swiglu.py new file mode 100644 index 00000000..a437d9a4 --- /dev/null +++ b/src/sparsevllm/triton_kernel/gate_up_swiglu.py @@ -0,0 +1,176 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from sparsevllm.triton_kernel.moe_config import MoeGemmConfig + + +_H20_DECODE_CONFIGS = { + (1, 2048, 512): MoeGemmConfig(16, 32, 64, 8, 4, 4), + (1, 2048, 256): MoeGemmConfig(16, 32, 64, 8, 4, 4), +} + + +def resolve_h20_gate_up_swiglu_config( + num_tokens: int, + hidden_size: int, + intermediate_size: int, +) -> MoeGemmConfig: + shape = (int(num_tokens), int(hidden_size), int(intermediate_size)) + try: + return _H20_DECODE_CONFIGS[shape] + except KeyError as error: + raise ValueError( + f"No H20 gate/up SwiGLU config for shape {shape}." + ) from error + + +@triton.jit +def _gate_up_swiglu_kernel( + input_ptr, + weight_ptr, + output_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + stride_am: tl.constexpr, + stride_ak: tl.constexpr, + stride_bn: tl.constexpr, + stride_bk: tl.constexpr, + stride_cm: tl.constexpr, + stride_cn: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m + pid_n = (pid % num_pid_in_group) // group_size_m + + m_offsets = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + n_offsets = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + k_offsets = tl.arange(0, BLOCK_SIZE_K) + input_ptrs = ( + input_ptr + + m_offsets[:, None] * stride_am + + k_offsets[None, :] * stride_ak + ) + gate_ptrs = ( + weight_ptr + + n_offsets[None, :] * stride_bn + + k_offsets[:, None] * stride_bk + ) + up_ptrs = gate_ptrs + N * stride_bn + gate_accumulator = tl.zeros( + (BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32 + ) + up_accumulator = tl.zeros( + (BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32 + ) + for k_start in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + remaining_k = K - k_start * BLOCK_SIZE_K + input_values = tl.load( + input_ptrs, + mask=(m_offsets[:, None] < M) + & (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_SIZE_K * stride_ak + gate_ptrs += BLOCK_SIZE_K * stride_bk + up_ptrs += BLOCK_SIZE_K * stride_bk + + 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] * stride_cm + + n_offsets[None, :] * stride_cn, + gate * up, + mask=(m_offsets[:, None] < M) & (n_offsets[None, :] < N), + ) + + +def gate_up_swiglu( + inputs: torch.Tensor, + weight: torch.Tensor, + config: MoeGemmConfig, + output: torch.Tensor | None = None, +) -> torch.Tensor: + if ( + inputs.ndim != 2 + or weight.ndim != 2 + or weight.shape[1] != inputs.shape[1] + ): + raise ValueError( + "gate_up_swiglu expects input [M, K] and weight [2N, K], " + f"got {tuple(inputs.shape)} and {tuple(weight.shape)}." + ) + if weight.shape[0] % 2: + raise ValueError( + f"gate_up_swiglu weight rows must be even, got {weight.shape[0]}." + ) + if inputs.dtype != torch.bfloat16 or weight.dtype != inputs.dtype: + raise TypeError("gate_up_swiglu requires matching BF16 inputs and weights.") + if not inputs.is_cuda or weight.device != inputs.device: + raise ValueError("gate_up_swiglu requires CUDA tensors on one device.") + if not inputs.is_contiguous() or not weight.is_contiguous(): + raise ValueError("gate_up_swiglu requires contiguous inputs and weights.") + + m, k = inputs.shape + n = weight.shape[0] // 2 + output = ( + torch.empty((m, n), dtype=inputs.dtype, device=inputs.device) + if output is None + else output + ) + if ( + output.shape != (m, n) + or output.dtype != inputs.dtype + or output.device != inputs.device + ): + raise ValueError( + f"gate_up_swiglu output must be {(m, n)} {inputs.dtype} on {inputs.device}." + ) + launch = config.as_triton_kwargs() + grid = ( + triton.cdiv(m, launch["BLOCK_SIZE_M"]) + * triton.cdiv(n, launch["BLOCK_SIZE_N"]), + ) + _gate_up_swiglu_kernel[grid]( + inputs, + weight, + output, + M=m, + N=n, + K=k, + stride_am=inputs.stride(0), + stride_ak=inputs.stride(1), + stride_bn=weight.stride(0), + stride_bk=weight.stride(1), + stride_cm=output.stride(0), + stride_cn=output.stride(1), + **launch, + ) + return output diff --git a/tests/test_gate_up_swiglu_operator.py b/tests/test_gate_up_swiglu_operator.py new file mode 100644 index 00000000..0d330031 --- /dev/null +++ b/tests/test_gate_up_swiglu_operator.py @@ -0,0 +1,85 @@ +import pytest +import torch + +from sparsevllm.operators.gate_up_swiglu import ( + GATE_UP_SWIGLU_REGISTRY, + GateUpSwiGLUOpSpec, + TorchGateUpSwiGLUProvider, +) +from sparsevllm.operators.registry import OpResolver +from sparsevllm.platforms import DeviceCaps, PlatformEnum + + +def _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 _caps(device_name="NVIDIA H20", capability=(9, 0)) -> DeviceCaps: + return DeviceCaps( + platform=PlatformEnum.CUDA, + device_type="cuda", + device_index=0, + device_name=device_name, + compute_capability=capability, + runtime_version="13.0", + supports_graph_capture=True, + supports_triton=True, + supports_bfloat16=True, + supports_native_fp8=True, + ) + + +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_h20_provider_requires_profiled_qwen36_shape(tp_size): + resolved = OpResolver(GATE_UP_SWIGLU_REGISTRY).resolve( + _spec(tp_size=tp_size), _caps() + ) + + assert resolved.provider.name == "h20_triton_decode" + + +@pytest.mark.parametrize( + "spec,caps", + [ + (_spec(weight_dtype=torch.float8_e4m3fn), _caps()), + (_spec(intermediate_size=768), _caps()), + (_spec(), _caps(device_name="NVIDIA H100 80GB HBM3")), + (_spec(), _caps(capability=(8, 9))), + ], +) +def test_unprofiled_shape_uses_torch_provider(spec, caps): + resolved = OpResolver(GATE_UP_SWIGLU_REGISTRY).resolve(spec, caps) + + assert resolved.provider.name == "torch" + + +def test_torch_provider_matches_gate_up_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 = TorchGateUpSwiGLUProvider().run( + _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) From 1f755ecb155a39be03b93433dc71a504b5704fb9 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 21:32:39 +0800 Subject: [PATCH 19/35] test: add shared expert tuner --- benchmark/tune_bf16_gate_up_swiglu.py | 103 ++++++++++++++++++++++++++ tests/test_gate_up_swiglu_kernel.py | 35 +++++++++ 2 files changed, 138 insertions(+) create mode 100644 benchmark/tune_bf16_gate_up_swiglu.py create mode 100644 tests/test_gate_up_swiglu_kernel.py diff --git a/benchmark/tune_bf16_gate_up_swiglu.py b/benchmark/tune_bf16_gate_up_swiglu.py new file mode 100644 index 00000000..c9b7cb0b --- /dev/null +++ b/benchmark/tune_bf16_gate_up_swiglu.py @@ -0,0 +1,103 @@ +import argparse +import json +from pathlib import Path + +import torch +import torch.nn.functional as F +import triton + +from sparsevllm.layers.activation import SiluAndMul +from sparsevllm.triton_kernel.gate_up_swiglu import gate_up_swiglu +from sparsevllm.triton_kernel.moe_config import MoeGemmConfig + + +def parse_args(): + parser = argparse.ArgumentParser(description="Tune BF16 gate/up GEMM + SwiGLU.") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--tokens", type=int, default=1) + parser.add_argument("--intermediate-size", type=int, default=512) + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--rep", type=int, default=300) + return parser.parse_args() + + +def candidates(): + for block_n in (16, 32, 64, 128): + for block_k, warps, stages in ( + (32, 4, 3), + (32, 4, 4), + (64, 4, 2), + (64, 4, 3), + (64, 4, 4), + (64, 8, 3), + ): + yield MoeGemmConfig(16, block_n, block_k, 8, warps, stages) + + +def main(): + args = parse_args() + if min(args.tokens, args.intermediate_size, args.warmup, args.rep) <= 0: + raise ValueError("All numeric arguments must be positive.") + torch.manual_seed(0) + inputs = torch.randn(args.tokens, 2048, dtype=torch.bfloat16, device="cuda") + weight = torch.randn( + 2 * args.intermediate_size, + 2048, + dtype=torch.bfloat16, + device="cuda", + ) + reference_projection = F.linear(inputs, weight) + gate, up = reference_projection.chunk(2, dim=-1) + reference = F.silu(gate.float()).mul(up.float()) + activation = SiluAndMul() + baseline = lambda: activation(F.linear(inputs, weight)) + baseline() + baseline_us = 1000 * triton.testing.do_bench( + baseline, warmup=args.warmup, rep=args.rep, return_mode="median" + ) + + records = [] + for config in candidates(): + output = torch.empty_like(reference_projection[:, : args.intermediate_size]) + launch = lambda: gate_up_swiglu(inputs, weight, config, output) + record = {**config.__dict__} + try: + actual = launch().float() + torch.cuda.synchronize() + record.update( + status="success", + max_abs_error=float((actual - reference).abs().max()), + relative_l2_error=float( + torch.linalg.vector_norm(actual - reference) + / torch.linalg.vector_norm(reference) + ), + latency_us=1000 + * triton.testing.do_bench( + launch, + warmup=args.warmup, + rep=args.rep, + return_mode="median", + ), + ) + except Exception as error: + record.update(status="invalid_config", error=f"{type(error).__name__}: {error}") + print(json.dumps(record, sort_keys=True), flush=True) + records.append(record) + + result = { + "device": torch.cuda.get_device_name(), + "torch_version": torch.__version__, + "triton_version": triton.__version__, + "tokens": args.tokens, + "intermediate_size": args.intermediate_size, + "warmup": args.warmup, + "rep": args.rep, + "baseline_us": baseline_us, + "records": records, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tests/test_gate_up_swiglu_kernel.py b/tests/test_gate_up_swiglu_kernel.py new file mode 100644 index 00000000..649abd6e --- /dev/null +++ b/tests/test_gate_up_swiglu_kernel.py @@ -0,0 +1,35 @@ +import pytest +import torch + +from sparsevllm.triton_kernel.gate_up_swiglu import ( + gate_up_swiglu, + resolve_h20_gate_up_swiglu_config, +) + + +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 = gate_up_swiglu( + inputs, + weight, + resolve_h20_gate_up_swiglu_config(1, 2048, intermediate_size), + ) + + torch.testing.assert_close(actual.float(), expected, rtol=0.02, atol=0.01) From 18e346e8f6853f2118534c8d089490844f07dbe6 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 21:32:43 +0800 Subject: [PATCH 20/35] perf: skip single-rank moe packing --- src/sparsevllm/models/qwen3_5_moe.py | 9 +++++---- tests/test_qwen35_mixed_runtime.py | 29 ++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/src/sparsevllm/models/qwen3_5_moe.py b/src/sparsevllm/models/qwen3_5_moe.py index 9a6cd885..aa2c871f 100644 --- a/src/sparsevllm/models/qwen3_5_moe.py +++ b/src/sparsevllm/models/qwen3_5_moe.py @@ -263,10 +263,11 @@ def _forward_chunk( 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) - routed_output, shared_output = self.parallel_context.world_all_reduce( - torch.stack((local_output, shared_output)) - ) - return gated_shared_add(routed_output, shared_output, shared_gate_logits) + 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: diff --git a/tests/test_qwen35_mixed_runtime.py b/tests/test_qwen35_mixed_runtime.py index 1b872458..b168f81c 100644 --- a/tests/test_qwen35_mixed_runtime.py +++ b/tests/test_qwen35_mixed_runtime.py @@ -202,6 +202,7 @@ def forward(self, *_): 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( @@ -218,6 +219,34 @@ def forward(self, *_): ) +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 From 674b080266c7005968ce2f1b41d79952f3827dd1 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 21:44:44 +0800 Subject: [PATCH 21/35] chore: remove moe tuning scripts --- benchmark/microbench.py | 4 - benchmark/tune_bf16_gate_up_swiglu.py | 103 ------------ benchmark/tune_bf16_moe.py | 218 -------------------------- benchmark/tune_fp8_moe.py | 146 ----------------- 4 files changed, 471 deletions(-) delete mode 100644 benchmark/tune_bf16_gate_up_swiglu.py delete mode 100644 benchmark/tune_bf16_moe.py delete mode 100644 benchmark/tune_fp8_moe.py diff --git a/benchmark/microbench.py b/benchmark/microbench.py index a1e0b3b0..e92191a4 100644 --- a/benchmark/microbench.py +++ b/benchmark/microbench.py @@ -647,7 +647,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) @@ -683,8 +682,6 @@ 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() @@ -725,7 +722,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/tune_bf16_gate_up_swiglu.py b/benchmark/tune_bf16_gate_up_swiglu.py deleted file mode 100644 index c9b7cb0b..00000000 --- a/benchmark/tune_bf16_gate_up_swiglu.py +++ /dev/null @@ -1,103 +0,0 @@ -import argparse -import json -from pathlib import Path - -import torch -import torch.nn.functional as F -import triton - -from sparsevllm.layers.activation import SiluAndMul -from sparsevllm.triton_kernel.gate_up_swiglu import gate_up_swiglu -from sparsevllm.triton_kernel.moe_config import MoeGemmConfig - - -def parse_args(): - parser = argparse.ArgumentParser(description="Tune BF16 gate/up GEMM + SwiGLU.") - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--tokens", type=int, default=1) - parser.add_argument("--intermediate-size", type=int, default=512) - parser.add_argument("--warmup", type=int, default=50) - parser.add_argument("--rep", type=int, default=300) - return parser.parse_args() - - -def candidates(): - for block_n in (16, 32, 64, 128): - for block_k, warps, stages in ( - (32, 4, 3), - (32, 4, 4), - (64, 4, 2), - (64, 4, 3), - (64, 4, 4), - (64, 8, 3), - ): - yield MoeGemmConfig(16, block_n, block_k, 8, warps, stages) - - -def main(): - args = parse_args() - if min(args.tokens, args.intermediate_size, args.warmup, args.rep) <= 0: - raise ValueError("All numeric arguments must be positive.") - torch.manual_seed(0) - inputs = torch.randn(args.tokens, 2048, dtype=torch.bfloat16, device="cuda") - weight = torch.randn( - 2 * args.intermediate_size, - 2048, - dtype=torch.bfloat16, - device="cuda", - ) - reference_projection = F.linear(inputs, weight) - gate, up = reference_projection.chunk(2, dim=-1) - reference = F.silu(gate.float()).mul(up.float()) - activation = SiluAndMul() - baseline = lambda: activation(F.linear(inputs, weight)) - baseline() - baseline_us = 1000 * triton.testing.do_bench( - baseline, warmup=args.warmup, rep=args.rep, return_mode="median" - ) - - records = [] - for config in candidates(): - output = torch.empty_like(reference_projection[:, : args.intermediate_size]) - launch = lambda: gate_up_swiglu(inputs, weight, config, output) - record = {**config.__dict__} - try: - actual = launch().float() - torch.cuda.synchronize() - record.update( - status="success", - max_abs_error=float((actual - reference).abs().max()), - relative_l2_error=float( - torch.linalg.vector_norm(actual - reference) - / torch.linalg.vector_norm(reference) - ), - latency_us=1000 - * triton.testing.do_bench( - launch, - warmup=args.warmup, - rep=args.rep, - return_mode="median", - ), - ) - except Exception as error: - record.update(status="invalid_config", error=f"{type(error).__name__}: {error}") - print(json.dumps(record, sort_keys=True), flush=True) - records.append(record) - - result = { - "device": torch.cuda.get_device_name(), - "torch_version": torch.__version__, - "triton_version": triton.__version__, - "tokens": args.tokens, - "intermediate_size": args.intermediate_size, - "warmup": args.warmup, - "rep": args.rep, - "baseline_us": baseline_us, - "records": records, - } - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(json.dumps(result, indent=2) + "\n") - - -if __name__ == "__main__": - main() diff --git a/benchmark/tune_bf16_moe.py b/benchmark/tune_bf16_moe.py deleted file mode 100644 index 04ffeefc..00000000 --- a/benchmark/tune_bf16_moe.py +++ /dev/null @@ -1,218 +0,0 @@ -import argparse -import json -from pathlib import Path - -import torch -import triton - -from sparsevllm.triton_kernel.moe import ( - _prepare_expert_assignment, - _routed_gate_up_swiglu, - _routed_gemm, -) -from sparsevllm.triton_kernel.moe_config import MoeGemmConfig -from sparsevllm.triton_kernel.silu_and_mul import silu_and_mul_fwd - - -def parse_args(): - parser = argparse.ArgumentParser(description="Tune Qwen3.6 BF16 routed GEMMs.") - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--intermediate-size", type=int, default=512) - parser.add_argument("--num-local-experts", type=int, default=256) - parser.add_argument("--local-assignments", type=int, default=8) - parser.add_argument("--warmup", type=int, default=50) - parser.add_argument("--rep", type=int, default=300) - return parser.parse_args() - - -def candidates(): - for block_n in (32, 64, 128): - for block_k, warps, stages in ( - (32, 4, 3), - (32, 4, 4), - (64, 4, 2), - (64, 4, 3), - (64, 4, 4), - (64, 8, 3), - ): - yield MoeGemmConfig(16, block_n, block_k, 8, warps, stages) - - -def benchmark(stage, launch, output, reference, config, args): - record = {"stage": stage, **config.__dict__} - try: - launch() - torch.cuda.synchronize() - actual = output[: args.local_assignments].float() - expected = reference[: args.local_assignments].float() - record.update( - status="success", - max_abs_error=float((actual - expected).abs().max()), - relative_l2_error=float( - torch.linalg.vector_norm(actual - expected) - / torch.linalg.vector_norm(expected) - ), - latency_us=1000 - * triton.testing.do_bench( - launch, - warmup=args.warmup, - rep=args.rep, - return_mode="median", - ), - ) - except Exception as error: - record.update(status="invalid_config", error=f"{type(error).__name__}: {error}") - print(json.dumps(record, sort_keys=True), flush=True) - return record - - -def main(): - args = parse_args() - if args.warmup <= 0 or args.rep <= 0: - raise ValueError("--warmup and --rep must be positive.") - if args.intermediate_size <= 0 or not 1 <= args.num_local_experts <= 256: - raise ValueError("--intermediate-size and --num-local-experts must be valid.") - if not 1 <= args.local_assignments <= min(8, args.num_local_experts): - raise ValueError("--local-assignments must fit the local expert range.") - if args.num_local_experts + 8 - args.local_assignments > 256: - raise ValueError("Remote assignments must fit the global expert range.") - torch.manual_seed(0) - intermediate_size = args.intermediate_size - num_weight_experts = args.local_assignments - hidden = torch.randn(1, 2048, device="cuda", dtype=torch.bfloat16) - w13 = torch.randn( - num_weight_experts, - 2 * intermediate_size, - 2048, - device="cuda", - dtype=torch.bfloat16, - ) - w2 = torch.randn( - num_weight_experts, - 2048, - intermediate_size, - device="cuda", - dtype=torch.bfloat16, - ) - expert_ids = list(range(args.local_assignments)) + list( - range( - args.num_local_experts, - args.num_local_experts + 8 - args.local_assignments, - ) - ) - topk_ids = torch.tensor([expert_ids], device="cuda", dtype=torch.int32) - topk_weights = torch.full((1, 8), 0.125, device="cuda", dtype=torch.bfloat16) - alignment = _prepare_expert_assignment( - topk_ids, - block_size=16, - num_experts=256, - local_expert_start=0, - local_expert_end=args.num_local_experts, - ) - reference_config = MoeGemmConfig(16, 128, 32, 8, 4, 4).as_triton_kwargs() - w13_reference = torch.empty( - 8, 2 * intermediate_size, device="cuda", dtype=torch.bfloat16 - ) - _routed_gemm( - hidden, - w13, - w13_reference, - topk_weights, - alignment, - input_top_k=8, - multiply_routing_weight=False, - launch_config=reference_config, - ) - activated = silu_and_mul_fwd(w13_reference.clone()) - w2_reference = torch.empty(8, 2048, device="cuda", dtype=torch.bfloat16) - _routed_gemm( - activated, - w2, - w2_reference, - topk_weights, - alignment, - input_top_k=1, - multiply_routing_weight=True, - launch_config=reference_config, - ) - torch.cuda.synchronize() - - records = [] - for config in candidates(): - launch_config = config.as_triton_kwargs() - w13_output = torch.empty_like(w13_reference) - records.append( - benchmark( - "w13", - lambda: _routed_gemm( - hidden, - w13, - w13_output, - topk_weights, - alignment, - input_top_k=8, - multiply_routing_weight=False, - launch_config=launch_config, - ), - w13_output, - w13_reference, - config, - args, - ) - ) - fused_output = torch.empty_like(activated) - records.append( - benchmark( - "gate_up_swiglu", - lambda: _routed_gate_up_swiglu( - hidden, - w13, - fused_output, - alignment, - input_top_k=8, - launch_config=launch_config, - ), - fused_output, - activated, - config, - args, - ) - ) - w2_output = torch.empty_like(w2_reference) - records.append( - benchmark( - "w2", - lambda: _routed_gemm( - activated, - w2, - w2_output, - topk_weights, - alignment, - input_top_k=1, - multiply_routing_weight=True, - launch_config=launch_config, - ), - w2_output, - w2_reference, - config, - args, - ) - ) - - result = { - "device": torch.cuda.get_device_name(), - "torch_version": torch.__version__, - "triton_version": triton.__version__, - "intermediate_size": intermediate_size, - "num_local_experts": args.num_local_experts, - "local_assignments": args.local_assignments, - "warmup": args.warmup, - "rep": args.rep, - "records": records, - } - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(json.dumps(result, indent=2) + "\n") - - -if __name__ == "__main__": - main() diff --git a/benchmark/tune_fp8_moe.py b/benchmark/tune_fp8_moe.py deleted file mode 100644 index df0a74a7..00000000 --- a/benchmark/tune_fp8_moe.py +++ /dev/null @@ -1,146 +0,0 @@ -import argparse -import json -from pathlib import Path - -import torch -import triton - -from sparsevllm.triton_kernel.moe import ( - _prepare_expert_assignment, - _routed_fp8_gemm, -) -from sparsevllm.triton_kernel.moe_config import MoeGemmConfig - - -def parse_args(): - parser = argparse.ArgumentParser(description="Tune Qwen3.6 EP2 FP8 routed GEMMs.") - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--local-assignments", type=int, default=4) - parser.add_argument("--warmup", type=int, default=50) - parser.add_argument("--rep", type=int, default=300) - return parser.parse_args() - - -def candidates(): - for block_n, swap_ab, stages in ( - (64, True, range(2, 6)), - (64, False, range(2, 5)), - (128, True, range(2, 5)), - (128, False, range(2, 5)), - (32, True, range(3, 5)), - ): - yield from ( - MoeGemmConfig(16, block_n, 128, 1, 4, stage, swap_ab) - for stage in stages - ) - - -def run_stage(stage, inputs, weights, scales, topk_weights, alignment, args): - multiply_routing_weight = stage == "w2" - input_top_k = 1 if multiply_routing_weight else 8 - reference_config = MoeGemmConfig(16, 128, 128, 1, 4, 3) - reference = torch.empty(8, weights.shape[1], device="cuda", dtype=torch.bfloat16) - _routed_fp8_gemm( - inputs, - weights, - scales, - reference, - topk_weights, - alignment, - input_top_k=input_top_k, - multiply_routing_weight=multiply_routing_weight, - config=reference_config, - ) - torch.cuda.synchronize() - records = [] - for config in candidates(): - output = torch.empty_like(reference) - - def launch(): - _routed_fp8_gemm( - inputs, - weights, - scales, - output, - topk_weights, - alignment, - input_top_k=input_top_k, - multiply_routing_weight=multiply_routing_weight, - config=config, - ) - - record = { - "stage": stage, - "local_assignments": args.local_assignments, - **config.__dict__, - } - try: - launch() - torch.cuda.synchronize() - count = args.local_assignments - actual, expected = output[:count].float(), reference[:count].float() - record.update( - status="success", - max_abs_error=float((actual - expected).abs().max()), - latency_us=1000 - * triton.testing.do_bench( - launch, - warmup=args.warmup, - rep=args.rep, - return_mode="median", - ), - ) - except Exception as error: - record.update(status="invalid_config", error=f"{type(error).__name__}: {error}") - records.append(record) - print(json.dumps(record, sort_keys=True), flush=True) - return records - - -def main(): - args = parse_args() - if args.warmup <= 0 or args.rep <= 0: - raise ValueError("--warmup and --rep must be positive.") - if not 1 <= args.local_assignments <= 8: - raise ValueError("--local-assignments must be in [1, 8].") - torch.manual_seed(0) - hidden = torch.randn(1, 2048, device="cuda", dtype=torch.bfloat16) - activated = torch.randn(8, 512, device="cuda", dtype=torch.bfloat16) - expert_ids = list(range(args.local_assignments)) + list( - range(128, 136 - args.local_assignments) - ) - topk_ids = torch.tensor([expert_ids], device="cuda", dtype=torch.int32) - topk_weights = torch.full((1, 8), 0.125, device="cuda", dtype=torch.bfloat16) - alignment = _prepare_expert_assignment( - topk_ids, - block_size=16, - num_experts=256, - local_expert_start=0, - local_expert_end=128, - ) - shapes = {"w13": (hidden, 1024, 2048), "w2": (activated, 2048, 512)} - records = [] - for stage, (inputs, output_size, input_size) in shapes.items(): - weights = torch.randn( - 8, output_size, input_size, device="cuda", dtype=torch.bfloat16 - ).to(torch.float8_e4m3fn) - scales = torch.ones( - 8, output_size // 128, input_size // 128, device="cuda", dtype=torch.bfloat16 - ) - records.extend( - run_stage( - stage, - inputs, - weights, - scales, - topk_weights, - alignment, - args, - ) - ) - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(json.dumps(records, indent=2) + "\n") - - -if __name__ == "__main__": - main() From 1b7a2cbb87b771ee00cda8903b27bc116d85754c Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 21:44:51 +0800 Subject: [PATCH 22/35] refactor: simplify h20 swiglu kernel --- src/sparsevllm/operators/gate_up_swiglu.py | 45 ++--- .../triton_kernel/gate_up_swiglu.py | 159 +++++------------- src/sparsevllm/triton_kernel/moe_config.py | 4 +- tests/test_gate_up_swiglu_kernel.py | 11 +- tests/test_gate_up_swiglu_operator.py | 10 +- 5 files changed, 65 insertions(+), 164 deletions(-) diff --git a/src/sparsevllm/operators/gate_up_swiglu.py b/src/sparsevllm/operators/gate_up_swiglu.py index b9deb222..c4ac8c6a 100644 --- a/src/sparsevllm/operators/gate_up_swiglu.py +++ b/src/sparsevllm/operators/gate_up_swiglu.py @@ -1,8 +1,6 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Protocol - import torch import torch.nn.functional as F @@ -11,12 +9,6 @@ from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum -class GateUpProjection(Protocol): - weight: torch.Tensor - - def __call__(self, inputs: torch.Tensor) -> torch.Tensor: ... - - @dataclass(frozen=True) class GateUpSwiGLUOpSpec: hidden_size: int @@ -33,8 +25,13 @@ def __post_init__(self) -> None: raise ValueError( "Gate/up SwiGLU intermediate size must be divisible by TP size." ) - if not self.activation_dtype.is_floating_point or not self.weight_dtype.is_floating_point: - raise TypeError("Gate/up SwiGLU activations and weights must be floating point.") + 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: @@ -45,7 +42,7 @@ def run( self, spec: GateUpSwiGLUOpSpec, inputs: torch.Tensor, - projection: GateUpProjection, + projection, ) -> torch.Tensor: raise NotImplementedError @@ -56,8 +53,8 @@ def run( @GATE_UP_SWIGLU_REGISTRY.register -class TorchGateUpSwiGLUProvider(GateUpSwiGLUProvider): - name = "torch" +class NativeGateUpSwiGLUProvider(GateUpSwiGLUProvider): + name = "native" priority = 0 @classmethod @@ -73,7 +70,7 @@ def run( self, spec: GateUpSwiGLUOpSpec, inputs: torch.Tensor, - projection: GateUpProjection, + projection, ) -> torch.Tensor: del spec gate, up = projection(inputs).chunk(2, dim=-1) @@ -81,7 +78,7 @@ def run( @GATE_UP_SWIGLU_REGISTRY.register -class H20GateUpSwiGLUProvider(TorchGateUpSwiGLUProvider): +class H20GateUpSwiGLUProvider(NativeGateUpSwiGLUProvider): name = "h20_triton_decode" priority = 20 @@ -125,23 +122,13 @@ def run( self, spec: GateUpSwiGLUOpSpec, inputs: torch.Tensor, - projection: GateUpProjection, + projection, ) -> torch.Tensor: if inputs.shape[0] != 1: return super().run(spec, inputs, projection) - from sparsevllm.triton_kernel.gate_up_swiglu import ( - gate_up_swiglu, - resolve_h20_gate_up_swiglu_config, - ) - - local_intermediate_size = spec.intermediate_size // spec.tp_size - return gate_up_swiglu( - inputs, - projection.weight, - resolve_h20_gate_up_swiglu_config( - inputs.shape[0], spec.hidden_size, local_intermediate_size - ), - ) + 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( diff --git a/src/sparsevllm/triton_kernel/gate_up_swiglu.py b/src/sparsevllm/triton_kernel/gate_up_swiglu.py index a437d9a4..8cb85370 100644 --- a/src/sparsevllm/triton_kernel/gate_up_swiglu.py +++ b/src/sparsevllm/triton_kernel/gate_up_swiglu.py @@ -4,83 +4,38 @@ import triton import triton.language as tl -from sparsevllm.triton_kernel.moe_config import MoeGemmConfig - _H20_DECODE_CONFIGS = { - (1, 2048, 512): MoeGemmConfig(16, 32, 64, 8, 4, 4), - (1, 2048, 256): MoeGemmConfig(16, 32, 64, 8, 4, 4), + (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), } -def resolve_h20_gate_up_swiglu_config( - num_tokens: int, - hidden_size: int, - intermediate_size: int, -) -> MoeGemmConfig: - shape = (int(num_tokens), int(hidden_size), int(intermediate_size)) - try: - return _H20_DECODE_CONFIGS[shape] - except KeyError as error: - raise ValueError( - f"No H20 gate/up SwiGLU config for shape {shape}." - ) from error - - @triton.jit def _gate_up_swiglu_kernel( input_ptr, weight_ptr, output_ptr, - M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, - stride_am: tl.constexpr, - stride_ak: tl.constexpr, - stride_bn: tl.constexpr, - stride_bk: tl.constexpr, - stride_cm: tl.constexpr, - stride_cn: tl.constexpr, - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, ): - pid = tl.program_id(0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m - pid_n = (pid % num_pid_in_group) // group_size_m + 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) - m_offsets = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - n_offsets = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - k_offsets = tl.arange(0, BLOCK_SIZE_K) - input_ptrs = ( - input_ptr - + m_offsets[:, None] * stride_am - + k_offsets[None, :] * stride_ak - ) - gate_ptrs = ( - weight_ptr - + n_offsets[None, :] * stride_bn - + k_offsets[:, None] * stride_bk - ) - up_ptrs = gate_ptrs + N * stride_bn - gate_accumulator = tl.zeros( - (BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32 - ) - up_accumulator = tl.zeros( - (BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32 - ) - for k_start in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - remaining_k = K - k_start * BLOCK_SIZE_K + 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] < M) + mask=(m_offsets[:, None] == 0) & (k_offsets[None, :] < remaining_k), other=0.0, ) @@ -95,82 +50,48 @@ def _gate_up_swiglu_kernel( input_values, tl.load(up_ptrs, mask=weight_mask, other=0.0), ) - input_ptrs += BLOCK_SIZE_K * stride_ak - gate_ptrs += BLOCK_SIZE_K * stride_bk - up_ptrs += BLOCK_SIZE_K * stride_bk + 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] * stride_cm - + n_offsets[None, :] * stride_cn, + output_ptr + m_offsets[:, None] * N + n_offsets[None, :], gate * up, - mask=(m_offsets[:, None] < M) & (n_offsets[None, :] < N), + mask=(m_offsets[:, None] == 0) & (n_offsets[None, :] < N), ) -def gate_up_swiglu( - inputs: torch.Tensor, - weight: torch.Tensor, - config: MoeGemmConfig, - output: torch.Tensor | None = None, -) -> torch.Tensor: - if ( - inputs.ndim != 2 - or weight.ndim != 2 - or weight.shape[1] != inputs.shape[1] - ): - raise ValueError( - "gate_up_swiglu expects input [M, K] and weight [2N, K], " - f"got {tuple(inputs.shape)} and {tuple(weight.shape)}." - ) - if weight.shape[0] % 2: +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( - f"gate_up_swiglu weight rows must be even, got {weight.shape[0]}." + "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("gate_up_swiglu requires matching BF16 inputs and weights.") + raise TypeError("h20_gate_up_swiglu requires matching BF16 tensors.") if not inputs.is_cuda or weight.device != inputs.device: - raise ValueError("gate_up_swiglu requires CUDA tensors on one 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("gate_up_swiglu requires contiguous inputs and weights.") + raise ValueError("h20_gate_up_swiglu requires contiguous tensors.") - m, k = inputs.shape - n = weight.shape[0] // 2 - output = ( - torch.empty((m, n), dtype=inputs.dtype, device=inputs.device) - if output is None - else output - ) - if ( - output.shape != (m, n) - or output.dtype != inputs.dtype - or output.device != inputs.device - ): - raise ValueError( - f"gate_up_swiglu output must be {(m, n)} {inputs.dtype} on {inputs.device}." - ) - launch = config.as_triton_kwargs() - grid = ( - triton.cdiv(m, launch["BLOCK_SIZE_M"]) - * triton.cdiv(n, launch["BLOCK_SIZE_N"]), - ) - _gate_up_swiglu_kernel[grid]( + 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, - M=m, - N=n, - K=k, - stride_am=inputs.stride(0), - stride_ak=inputs.stride(1), - stride_bn=weight.stride(0), - stride_bk=weight.stride(1), - stride_cm=output.stride(0), - stride_cn=output.stride(1), - **launch, + 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_config.py b/src/sparsevllm/triton_kernel/moe_config.py index 521b6ff6..5b196728 100644 --- a/src/sparsevllm/triton_kernel/moe_config.py +++ b/src/sparsevllm/triton_kernel/moe_config.py @@ -204,8 +204,8 @@ def _stage_table( _FP8_N128_SWAP_S4 = MoeGemmConfig(16, 128, 128, 1, 4, 4, True) -# Qwen3.6-35B-A3B block-FP8 decode profiles tuned offline on H100. Larger -# token buckets retain the explicit generic configuration until profiled. +# Qwen3.6-35B-A3B block-FP8 decode profiles. Unprofiled token buckets retain +# the explicit generic configuration. _TUNED_FP8_ROUTED_CONFIGS = { MoeGemmShape( "H20", diff --git a/tests/test_gate_up_swiglu_kernel.py b/tests/test_gate_up_swiglu_kernel.py index 649abd6e..e9e0f926 100644 --- a/tests/test_gate_up_swiglu_kernel.py +++ b/tests/test_gate_up_swiglu_kernel.py @@ -1,10 +1,7 @@ import pytest import torch -from sparsevllm.triton_kernel.gate_up_swiglu import ( - gate_up_swiglu, - resolve_h20_gate_up_swiglu_config, -) +from sparsevllm.triton_kernel.gate_up_swiglu import h20_gate_up_swiglu def _is_h20() -> bool: @@ -26,10 +23,6 @@ def test_h20_gate_up_swiglu_matches_torch(intermediate_size): gate, up = projected.chunk(2, dim=-1) expected = torch.nn.functional.silu(gate.float()) * up.float() - actual = gate_up_swiglu( - inputs, - weight, - resolve_h20_gate_up_swiglu_config(1, 2048, intermediate_size), - ) + actual = h20_gate_up_swiglu(inputs, weight) torch.testing.assert_close(actual.float(), expected, rtol=0.02, atol=0.01) diff --git a/tests/test_gate_up_swiglu_operator.py b/tests/test_gate_up_swiglu_operator.py index 0d330031..3d0cef5d 100644 --- a/tests/test_gate_up_swiglu_operator.py +++ b/tests/test_gate_up_swiglu_operator.py @@ -4,7 +4,7 @@ from sparsevllm.operators.gate_up_swiglu import ( GATE_UP_SWIGLU_REGISTRY, GateUpSwiGLUOpSpec, - TorchGateUpSwiGLUProvider, + NativeGateUpSwiGLUProvider, ) from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum @@ -56,13 +56,13 @@ def test_h20_provider_requires_profiled_qwen36_shape(tp_size): (_spec(), _caps(capability=(8, 9))), ], ) -def test_unprofiled_shape_uses_torch_provider(spec, caps): +def test_unprofiled_shape_uses_native_provider(spec, caps): resolved = OpResolver(GATE_UP_SWIGLU_REGISTRY).resolve(spec, caps) - assert resolved.provider.name == "torch" + assert resolved.provider.name == "native" -def test_torch_provider_matches_gate_up_swiglu_semantics(): +def test_native_provider_matches_gate_up_swiglu_semantics(): torch.manual_seed(0) inputs = torch.randn(3, 8) projection = torch.nn.Linear(8, 12, bias=False) @@ -70,7 +70,7 @@ def test_torch_provider_matches_gate_up_swiglu_semantics(): projected = projection(inputs) gate, up = projected.chunk(2, dim=-1) expected = torch.nn.functional.silu(gate) * up - actual = TorchGateUpSwiGLUProvider().run( + actual = NativeGateUpSwiGLUProvider().run( _spec( hidden_size=8, intermediate_size=6, From 2488533561b6d803051970057c7b5eba965857df Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 21:44:56 +0800 Subject: [PATCH 23/35] refactor: isolate gated shared add --- src/sparsevllm/models/qwen3_5_moe.py | 2 +- src/sparsevllm/operators/gated_shared_add.py | 44 +++++++++++++++++++ .../triton_kernel/qwen3_5/gated_shared_add.py | 31 +------------ tests/test_qwen35_gated_shared_add.py | 2 +- 4 files changed, 47 insertions(+), 32 deletions(-) create mode 100644 src/sparsevllm/operators/gated_shared_add.py diff --git a/src/sparsevllm/models/qwen3_5_moe.py b/src/sparsevllm/models/qwen3_5_moe.py index aa2c871f..f0edae0e 100644 --- a/src/sparsevllm/models/qwen3_5_moe.py +++ b/src/sparsevllm/models/qwen3_5_moe.py @@ -19,13 +19,13 @@ 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.triton_kernel.qwen3_5.gated_shared_add import gated_shared_add from sparsevllm.utils.log import logger from sparsevllm.utils.weight_target import WeightTarget 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/triton_kernel/qwen3_5/gated_shared_add.py b/src/sparsevllm/triton_kernel/qwen3_5/gated_shared_add.py index 0cc67525..1fe35f8e 100644 --- a/src/sparsevllm/triton_kernel/qwen3_5/gated_shared_add.py +++ b/src/sparsevllm/triton_kernel/qwen3_5/gated_shared_add.py @@ -30,40 +30,11 @@ def _gated_shared_add_kernel( ) -def gated_shared_add( +def triton_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.") - output = torch.empty_like(routed) hidden_size = int(routed.shape[1]) block_size = 512 diff --git a/tests/test_qwen35_gated_shared_add.py b/tests/test_qwen35_gated_shared_add.py index 635ece37..71dda157 100644 --- a/tests/test_qwen35_gated_shared_add.py +++ b/tests/test_qwen35_gated_shared_add.py @@ -1,7 +1,7 @@ import pytest import torch -from sparsevllm.triton_kernel.qwen3_5.gated_shared_add import gated_shared_add +from sparsevllm.operators.gated_shared_add import gated_shared_add pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") From 32f971fb390586d1fe966f771c2bbd349b670f31 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 22:34:02 +0800 Subject: [PATCH 24/35] perf: tune qwen3.6 moe buckets --- src/sparsevllm/operators/moe.py | 9 +- src/sparsevllm/triton_kernel/moe_config.py | 140 ++++++++++++++++++--- 2 files changed, 128 insertions(+), 21 deletions(-) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index e1a7c12d..5ab1ae7e 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -357,7 +357,12 @@ class TritonHopperFusedMoeProvider(MoeProvider): priority = 20 gate_up_order = "gate_up" PROFILED_DEVICE_NAME = "NVIDIA H100 80GB HBM3" - PROFILED_SHAPES = ((128, 64, 2048, 384, 8, 2, 2),) + 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: @@ -620,7 +625,7 @@ class H20Qwen36HybridFp8MoeProvider(HopperQwen36HybridFp8MoeProvider): name = "h20_qwen36_hybrid_fp8" priority = 111 PROFILED_DEVICE_NAME = "NVIDIA H20" - TRITON_MAX_TOKENS_BY_EP_SIZE = {1: 1, 2: 1} + TRITON_MAX_TOKENS_BY_EP_SIZE = {1: 8, 2: 1} def resolve_moe_provider( diff --git a/src/sparsevllm/triton_kernel/moe_config.py b/src/sparsevllm/triton_kernel/moe_config.py index 5b196728..d19837ca 100644 --- a/src/sparsevllm/triton_kernel/moe_config.py +++ b/src/sparsevllm/triton_kernel/moe_config.py @@ -94,6 +94,8 @@ def _heuristic_config( _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( @@ -112,16 +114,16 @@ def _stage_table( # than by model name. _TUNED_CONFIGS = { MoeGemmShape("H20", (9, 0), torch.bfloat16, 8, 256, 2048, 512): { - "w13": {1: _G}, - "w2": {1: _H}, + "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}, - "w2": {1: _H}, + "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}, - "w2": {1: _I}, + "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), @@ -147,6 +149,18 @@ 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}, + }, } @@ -161,7 +175,7 @@ def _stage_table( 256, 2048, 512, - ): {1: _G}, + ): {1: _G, 2: _G, 4: _H, 8: _G}, MoeGemmShape( "H20", (9, 0), @@ -170,7 +184,7 @@ def _stage_table( 256, 2048, 256, - ): {1: _G}, + ): {1: _G, 2: _G, 4: _G, 8: _H}, MoeGemmShape( "H20", (9, 0), @@ -179,7 +193,34 @@ def _stage_table( 128, 2048, 512, - ): {1: _G}, + ): {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), @@ -198,9 +239,11 @@ 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) @@ -216,8 +259,35 @@ def _stage_table( 2048, 512, ): { - "w13": {1: _FP8_N64_SWAP_S4}, - "w2": {1: _FP8_N128_SWAP_S4}, + "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", @@ -228,8 +298,13 @@ def _stage_table( 2048, 512, ): { - "w13": {1: _FP8_N64_SWAP_S4}, - "w2": {1: _FP8_N64_SWAP}, + "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", @@ -241,17 +316,39 @@ def _stage_table( 512, ): { "w13": { - 1: _FP8_N64_SWAP, + 1: _FP8_N64_SWAP_S5, 2: _FP8_N64_SWAP_S4, 4: _FP8_N64_SWAP, - 8: _FP8_N128, + 8: _FP8_N64_SWAP_S4, }, "w2": { 1: _FP8_N64_SWAP_S4, - 2: _FP8_N64_SWAP, - 4: _FP8_N128, + 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", @@ -262,8 +359,13 @@ def _stage_table( 2048, 512, ): { - "w13": {1: _FP8_N64_SWAP_S5}, - "w2": {1: _FP8_N64_SWAP_S4}, + "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, + }, }, } From 94a5b85c2e81f15d9d75e05108df635f178490df Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 22:34:05 +0800 Subject: [PATCH 25/35] test: cover qwen3.6 moe buckets --- tests/test_moe_config.py | 146 ++++++++++++++++++++++--------- tests/test_operator_providers.py | 6 +- 2 files changed, 110 insertions(+), 42 deletions(-) diff --git a/tests/test_moe_config.py b/tests/test_moe_config.py index 722e9938..10a93975 100644 --- a/tests/test_moe_config.py +++ b/tests/test_moe_config.py @@ -106,10 +106,19 @@ def test_h20_qwen3_moe_config_is_shape_and_stage_aware(): assert large.block_m == 64 -def test_h20_qwen36_decode_uses_profiled_bf16_configs(): +@pytest.mark.parametrize( + ("tokens", "expected"), + [ + (1, ((32, 64, 4, 4), (32, 64, 4, 4), (32, 64, 4, 4))), + (2, ((32, 64, 4, 4), (32, 64, 4, 4), (64, 64, 4, 3))), + (4, ((64, 64, 4, 3), (32, 64, 4, 3), (128, 64, 4, 2))), + (8, ((128, 64, 8, 3), (32, 64, 4, 4), (64, 64, 4, 4))), + ], +) +def test_h20_qwen36_decode_uses_profiled_bf16_configs(tokens, expected): common = dict( dtype=torch.bfloat16, - num_tokens=1, + num_tokens=tokens, top_k=8, num_local_experts=256, hidden_size=2048, @@ -121,15 +130,40 @@ def test_h20_qwen36_decode_uses_profiled_bf16_configs(): fused = resolve_moe_gemm_config(**common, stage="gate_up_swiglu") w2 = resolve_moe_gemm_config(**common, stage="w2") - assert (w13.block_n, w13.block_k, w13.num_stages) == (32, 64, 4) - assert fused == w13 - assert (w2.block_n, w2.block_k, w2.num_stages) == (32, 64, 3) + assert tuple( + (config.block_n, config.block_k, config.num_warps, config.num_stages) + for config in (w13, fused, w2) + ) == expected - unprofiled = resolve_moe_gemm_config( - **{**common, "num_tokens": 2}, - stage="w13", + +@pytest.mark.parametrize( + ("device_name", "num_local_experts", "intermediate_size", "expected"), + [ + ("NVIDIA H100 80GB HBM3", 256, 512, ((64, 3), (32, 3), (32, 3))), + ("NVIDIA H100 80GB HBM3", 256, 256, ((32, 3), (32, 3), (32, 3))), + ("NVIDIA H100 80GB HBM3", 128, 512, ((64, 3), (32, 3), (32, 3))), + ("NVIDIA H20", 256, 256, ((64, 3), (32, 3), (128, 2))), + ("NVIDIA H20", 128, 512, ((128, 3), (32, 3), (64, 4))), + ], +) +def test_qwen36_parallel_bf16_profiles_cover_bs8( + device_name, num_local_experts, intermediate_size, expected +): + common = dict( + dtype=torch.bfloat16, + num_tokens=8, + top_k=8, + num_local_experts=num_local_experts, + hidden_size=2048, + intermediate_size=intermediate_size, + device_name=device_name, + device_capability=(9, 0), ) - assert (unprofiled.block_n, unprofiled.block_k) == (128, 32) + configs = tuple( + resolve_moe_gemm_config(**common, stage=stage) + for stage in ("w13", "gate_up_swiglu", "w2") + ) + assert tuple((config.block_n, config.num_stages) for config in configs) == expected @pytest.mark.parametrize( @@ -213,11 +247,11 @@ def test_h100_profile_switches_to_large_token_config(): @pytest.mark.parametrize( ("stage", "tokens", "block_n", "num_stages", "swap_ab"), [ - ("w13", 1, 64, 3, True), + ("w13", 1, 64, 5, True), ("w13", 2, 64, 4, True), - ("w13", 8, 128, 3, False), + ("w13", 8, 64, 4, True), ("w2", 1, 64, 4, True), - ("w2", 4, 128, 3, False), + ("w2", 4, 64, 2, True), ("w2", 8, 64, 3, True), ], ) @@ -240,9 +274,20 @@ def test_h100_qwen36_fp8_routed_config(stage, tokens, block_n, num_stages, swap_ ) -def test_h100_qwen36_fp8_ep2_uses_profiled_configs(): +@pytest.mark.parametrize( + ("tokens", "expected_w13", "expected_w2"), + [ + (1, (64, 4), (64, 4)), + (2, (64, 4), (64, 3)), + (4, (64, 4), (64, 2)), + (8, (64, 4), (64, 2)), + ], +) +def test_h100_qwen36_fp8_ep2_uses_profiled_configs( + tokens, expected_w13, expected_w2 +): common = dict( - num_tokens=1, + num_tokens=tokens, top_k=8, num_local_experts=128, hidden_size=2048, @@ -253,38 +298,61 @@ def test_h100_qwen36_fp8_ep2_uses_profiled_configs(): w13 = resolve_fp8_routed_gemm_config(**common, stage="w13") w2 = resolve_fp8_routed_gemm_config(**common, stage="w2") - assert (w13.block_n, w13.num_stages, w13.swap_ab) == (64, 5, True) - assert (w2.block_n, w2.num_stages, w2.swap_ab) == (64, 4, True) + assert (w13.block_n, w13.num_stages) == expected_w13 + assert (w2.block_n, w2.num_stages) == expected_w2 + assert w13.swap_ab and w2.swap_ab @pytest.mark.parametrize( - ("local_experts", "stage", "block_n", "num_stages"), + ("local_experts", "intermediate_size", "expected"), [ - (256, "w13", 64, 4), - (256, "w2", 128, 4), - (128, "w13", 64, 4), - (128, "w2", 64, 3), + ( + 256, + 512, + { + "w13": ((64, 4), (128, 4), (64, 4), (64, 3)), + "w2": ((64, 3), (64, 2), (64, 2), (128, 2)), + }, + ), + ( + 256, + 256, + { + "w13": ((64, 4),) * 4, + "w2": ((64, 3), (64, 2), (64, 2), (128, 2)), + }, + ), + ( + 128, + 512, + { + "w13": ((64, 4),) * 4, + "w2": ((64, 3), (128, 4), (64, 2), (64, 2)), + }, + ), ], ) -def test_h20_qwen36_fp8_uses_profiled_decode_configs( - local_experts, stage, block_n, num_stages +def test_h20_qwen36_fp8_profiles_cover_decode_buckets( + local_experts, intermediate_size, expected ): - config = resolve_fp8_routed_gemm_config( - num_tokens=1, - top_k=8, - num_local_experts=local_experts, - hidden_size=2048, - intermediate_size=512, - stage=stage, - device_name="NVIDIA H20", - device_capability=(9, 0), - ) - - assert (config.block_n, config.num_stages, config.swap_ab) == ( - block_n, - num_stages, - True, - ) + 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="NVIDIA H20", + 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.swap_ab for config in configs) def test_fp8_routed_unknown_shape_uses_explicit_default(): diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 47c043b5..8e61670d 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -370,7 +370,7 @@ def test_hopper_fused_moe_uses_profiled_tp_ep_shape(): ("tp_size", "ep_size", "intermediate_size", "num_local_experts"), [(1, 1, 512, 256), (2, 1, 256, 256), (1, 2, 512, 128)], ) -def test_h20_qwen36_bf16_moe_uses_profiled_provider( +def test_qwen36_bf16_moe_uses_profiled_hopper_provider( tp_size, ep_size, intermediate_size, @@ -401,7 +401,7 @@ def test_h20_qwen36_bf16_moe_uses_profiled_provider( (9, 0), native_fp8=False, device_name="NVIDIA H100 80GB HBM3" ), ) - assert h100.provider.name == "triton" + assert h100.provider.name == "triton_hopper_fused" @pytest.mark.parametrize( @@ -669,7 +669,7 @@ def test_h20_qwen36_hybrid_moe_uses_profiled_provider(): def test_h20_qwen36_hybrid_moe_limits_triton_to_profiled_token_count(): - assert H20Qwen36HybridFp8MoeProvider.TRITON_MAX_TOKENS_BY_EP_SIZE == {1: 1, 2: 1} + assert H20Qwen36HybridFp8MoeProvider.TRITON_MAX_TOKENS_BY_EP_SIZE == {1: 8, 2: 1} def test_qwen36_hybrid_moe_dispatches_by_token_bucket(): From 8909e59b6789d2ac195f73ee735321dec843534c Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 22:37:39 +0800 Subject: [PATCH 26/35] test: verify all qwen3.6 moe buckets --- tests/test_moe_config.py | 263 ++++++++++++++++++--------------------- 1 file changed, 122 insertions(+), 141 deletions(-) diff --git a/tests/test_moe_config.py b/tests/test_moe_config.py index 10a93975..a60cf648 100644 --- a/tests/test_moe_config.py +++ b/tests/test_moe_config.py @@ -106,90 +106,94 @@ def test_h20_qwen3_moe_config_is_shape_and_stage_aware(): assert large.block_m == 64 -@pytest.mark.parametrize( - ("tokens", "expected"), - [ - (1, ((32, 64, 4, 4), (32, 64, 4, 4), (32, 64, 4, 4))), - (2, ((32, 64, 4, 4), (32, 64, 4, 4), (64, 64, 4, 3))), - (4, ((64, 64, 4, 3), (32, 64, 4, 3), (128, 64, 4, 2))), - (8, ((128, 64, 8, 3), (32, 64, 4, 4), (64, 64, 4, 4))), - ], -) -def test_h20_qwen36_decode_uses_profiled_bf16_configs(tokens, expected): - common = dict( - dtype=torch.bfloat16, - num_tokens=tokens, - top_k=8, - num_local_experts=256, - hidden_size=2048, - intermediate_size=512, - device_name="NVIDIA H20", - device_capability=(9, 0), - ) - w13 = resolve_moe_gemm_config(**common, stage="w13") - fused = resolve_moe_gemm_config(**common, stage="gate_up_swiglu") - w2 = resolve_moe_gemm_config(**common, stage="w2") - - assert tuple( - (config.block_n, config.block_k, config.num_warps, config.num_stages) - for config in (w13, fused, w2) - ) == expected - - @pytest.mark.parametrize( ("device_name", "num_local_experts", "intermediate_size", "expected"), [ - ("NVIDIA H100 80GB HBM3", 256, 512, ((64, 3), (32, 3), (32, 3))), - ("NVIDIA H100 80GB HBM3", 256, 256, ((32, 3), (32, 3), (32, 3))), - ("NVIDIA H100 80GB HBM3", 128, 512, ((64, 3), (32, 3), (32, 3))), - ("NVIDIA H20", 256, 256, ((64, 3), (32, 3), (128, 2))), - ("NVIDIA H20", 128, 512, ((128, 3), (32, 3), (64, 4))), + ( + "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_parallel_bf16_profiles_cover_bs8( +def test_qwen36_bf16_profiles_cover_decode_buckets( device_name, num_local_experts, intermediate_size, expected ): - common = dict( - dtype=torch.bfloat16, - num_tokens=8, - top_k=8, - num_local_experts=num_local_experts, - hidden_size=2048, - intermediate_size=intermediate_size, - device_name=device_name, - device_capability=(9, 0), - ) - configs = tuple( - resolve_moe_gemm_config(**common, stage=stage) - for stage in ("w13", "gate_up_swiglu", "w2") - ) - assert tuple((config.block_n, config.num_stages) for config in configs) == expected - - -@pytest.mark.parametrize( - ("num_local_experts", "intermediate_size", "expected_w2"), - [(256, 256, (32, 64, 3)), (128, 512, (64, 64, 4))], -) -def test_h20_qwen36_parallel_decode_uses_profiled_bf16_configs( - num_local_experts, - intermediate_size, - expected_w2, -): - common = dict( - dtype=torch.bfloat16, - num_tokens=1, - top_k=8, - num_local_experts=num_local_experts, - hidden_size=2048, - intermediate_size=intermediate_size, - device_name="NVIDIA H20", - device_capability=(9, 0), - ) - - fused = resolve_moe_gemm_config(**common, stage="gate_up_swiglu") - w2 = resolve_moe_gemm_config(**common, stage="w2") - assert (fused.block_n, fused.block_k, fused.num_stages) == (32, 64, 4) - assert (w2.block_n, w2.block_k, w2.num_stages) == expected_w2 + 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(): @@ -245,68 +249,10 @@ def test_h100_profile_switches_to_large_token_config(): @pytest.mark.parametrize( - ("stage", "tokens", "block_n", "num_stages", "swap_ab"), - [ - ("w13", 1, 64, 5, True), - ("w13", 2, 64, 4, True), - ("w13", 8, 64, 4, True), - ("w2", 1, 64, 4, True), - ("w2", 4, 64, 2, True), - ("w2", 8, 64, 3, True), - ], -) -def test_h100_qwen36_fp8_routed_config(stage, tokens, block_n, num_stages, swap_ab): - config = resolve_fp8_routed_gemm_config( - num_tokens=tokens, - top_k=8, - num_local_experts=256, - hidden_size=2048, - intermediate_size=512, - stage=stage, - device_name="NVIDIA H100 80GB HBM3", - device_capability=(9, 0), - ) - - assert (config.block_n, config.num_stages, config.swap_ab) == ( - block_n, - num_stages, - swap_ab, - ) - - -@pytest.mark.parametrize( - ("tokens", "expected_w13", "expected_w2"), - [ - (1, (64, 4), (64, 4)), - (2, (64, 4), (64, 3)), - (4, (64, 4), (64, 2)), - (8, (64, 4), (64, 2)), - ], -) -def test_h100_qwen36_fp8_ep2_uses_profiled_configs( - tokens, expected_w13, expected_w2 -): - common = dict( - num_tokens=tokens, - top_k=8, - num_local_experts=128, - hidden_size=2048, - intermediate_size=512, - device_name="NVIDIA H100 80GB HBM3", - device_capability=(9, 0), - ) - w13 = resolve_fp8_routed_gemm_config(**common, stage="w13") - w2 = resolve_fp8_routed_gemm_config(**common, stage="w2") - - assert (w13.block_n, w13.num_stages) == expected_w13 - assert (w2.block_n, w2.num_stages) == expected_w2 - assert w13.swap_ab and w2.swap_ab - - -@pytest.mark.parametrize( - ("local_experts", "intermediate_size", "expected"), + ("device_name", "local_experts", "intermediate_size", "expected"), [ ( + "NVIDIA H20", 256, 512, { @@ -315,6 +261,7 @@ def test_h100_qwen36_fp8_ep2_uses_profiled_configs( }, ), ( + "NVIDIA H20", 256, 256, { @@ -323,6 +270,7 @@ def test_h100_qwen36_fp8_ep2_uses_profiled_configs( }, ), ( + "NVIDIA H20", 128, 512, { @@ -330,10 +278,37 @@ def test_h100_qwen36_fp8_ep2_uses_profiled_configs( "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_h20_qwen36_fp8_profiles_cover_decode_buckets( - local_experts, intermediate_size, expected +def test_qwen36_fp8_profiles_cover_decode_buckets( + device_name, local_experts, intermediate_size, expected ): for stage, stage_expected in expected.items(): configs = tuple( @@ -344,7 +319,7 @@ def test_h20_qwen36_fp8_profiles_cover_decode_buckets( hidden_size=2048, intermediate_size=intermediate_size, stage=stage, - device_name="NVIDIA H20", + device_name=device_name, device_capability=(9, 0), ) for tokens in (1, 2, 4, 8) @@ -352,7 +327,13 @@ def test_h20_qwen36_fp8_profiles_cover_decode_buckets( assert tuple( (config.block_n, config.num_stages) for config in configs ) == stage_expected - assert all(config.swap_ab for config in configs) + 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(): From 45b808a2870b9e69bc7b07c5d6f140cc6fdd6767 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Tue, 11 Aug 2026 23:28:25 +0800 Subject: [PATCH 27/35] test: lock pure tp moe provider --- tests/test_operator_providers.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 8e61670d..9d34ab44 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -592,6 +592,25 @@ def test_qwen36_hybrid_moe_uses_profiled_single_gpu_shape_on_h100(): 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"), [ From d5ed96635bbb6c8b21c019ab0c1a4647d672f86a Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 14:23:33 +0800 Subject: [PATCH 28/35] fix: remove ambiguous e2e metric --- benchmark/microbench.py | 13 +++---------- tests/test_microbench_artifacts.py | 8 ++------ 2 files changed, 5 insertions(+), 16 deletions(-) diff --git a/benchmark/microbench.py b/benchmark/microbench.py index e92191a4..765aeaf7 100644 --- a/benchmark/microbench.py +++ b/benchmark/microbench.py @@ -286,8 +286,6 @@ def _artifact_records(args, rows: list[dict[str, Any]]) -> list[dict[str, Any]]: record.setdefault("itl_ms", row["itl"]) if "mem" in row: record.setdefault("peak_memory_gb", row["mem"]) - if "duration_s" in row: - record.setdefault("e2e_latency_s", row["duration_s"]) records.append(record) return records @@ -347,19 +345,18 @@ def _write_output_dir(args, rows: list[dict[str, Any]]) -> None: f"- Batch sizes: `{args.batch_sizes}`", f"- Output length: `{args.output_len}`", "", - "| Method | Prompt tokens | Batch | Status | E2E s | TTFT s | Prefill tok/s | Decode tok/s | Peak GB | Decode speedup |", - "| --- | ---: | ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: |", + "| Method | Prompt tokens | Batch | Status | TTFT s | Prefill tok/s | Decode tok/s | Peak GB | Decode speedup |", + "| --- | ---: | ---: | --- | ---: | ---: | ---: | ---: | ---: |", ] for record in records: ok = record["status"] == "success" speedup = record.get("speedup_vs_vanilla_decode") report_lines.append( - "| {method} | {prompt} | {batch} | {status} | {e2e} | {ttft} | {prefill} | {decode} | {mem} | {speedup} |".format( + "| {method} | {prompt} | {batch} | {status} | {ttft} | {prefill} | {decode} | {mem} | {speedup} |".format( method=record.get("method", ""), prompt=record.get("prompt_tokens", ""), batch=record.get("batch_size", ""), status=record["status"], - e2e=f"{record.get('e2e_latency_s', 0.0):.3f}" if ok else "", ttft=f"{record.get('ttft_s', 0.0):.3f}" if ok else "", prefill=f"{record.get('prefill_tok_s', 0.0):.1f}" if ok else "", decode=f"{record.get('decode_tok_s', 0.0):.1f}" if ok else "", @@ -683,9 +680,6 @@ def add_wave(max_new_requests: int): ) 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) @@ -753,7 +747,6 @@ def add_wave(max_new_requests: int): "prefill_tp": prefill_tp, "decode_tp": decode_tp, "ttft": ttft, - "duration_s": duration, "itl": avg_itl, "avg_bs": avg_active_bs, "mem": peak_mem, diff --git a/tests/test_microbench_artifacts.py b/tests/test_microbench_artifacts.py index f4cbdf27..3bdc0c99 100644 --- a/tests/test_microbench_artifacts.py +++ b/tests/test_microbench_artifacts.py @@ -165,7 +165,7 @@ def test_benchmark_sparse_method_preserves_graph_enabling_legacy_alias(method): @pytest.mark.parametrize("enabled", [False, True]) -def test_artifact_records_include_step_timing_mode_and_e2e_latency(enabled): +def test_artifact_records_include_step_timing_mode(enabled): args = SimpleNamespace( output_len=8, temperature=0.0, @@ -173,13 +173,9 @@ def test_artifact_records_include_step_timing_mode_and_e2e_latency(enabled): synchronize_step_timing=enabled, ) - records = _artifact_records( - args, - [{"status": "SUCCESS", "length": 16, "duration_s": 1.25}], - ) + records = _artifact_records(args, [{"status": "SUCCESS", "length": 16}]) assert records[0]["synchronize_step_timing"] is enabled - assert records[0]["e2e_latency_s"] == 1.25 def test_output_metadata_records_step_timing_mode(tmp_path, monkeypatch): From 8daebec844b06a0834c0fae8d63893dca97e076c Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 14:23:55 +0800 Subject: [PATCH 29/35] refactor: centralize model runtime rules --- src/sparsevllm/configs/common.py | 4 - src/sparsevllm/configs/delta.py | 45 +- src/sparsevllm/configs/model.py | 1060 ++--------------- src/sparsevllm/configs/platform.py | 11 +- src/sparsevllm/configs/prefix_cache.py | 21 +- src/sparsevllm/configs/runtime.py | 40 +- src/sparsevllm/configs/sparse.py | 9 +- src/sparsevllm/distributed/__init__.py | 6 + .../distributed/parallel_context.py | 42 +- src/sparsevllm/distributed/sharding.py | 62 + src/sparsevllm/distributed/topology.py | 73 ++ src/sparsevllm/engine/llm_engine.py | 11 +- src/sparsevllm/engine/model_runner.py | 71 +- src/sparsevllm/method_registry.py | 193 +-- src/sparsevllm/models/checkpoint.py | 279 +++++ src/sparsevllm/models/layout.py | 219 ++++ src/sparsevllm/models/spec.py | 184 +++ src/sparsevllm/quantization/__init__.py | 2 + src/sparsevllm/quantization/config.py | 124 ++ src/sparsevllm/utils/config.py | 7 + tests/test_minimax_m2_config.py | 48 +- tests/test_parallel_context.py | 54 +- tests/test_prefill_schedule_policy.py | 2 +- tests/test_qwen35_mixed_runtime.py | 42 +- tests/test_qwen3_moe_compatibility.py | 41 +- 25 files changed, 1338 insertions(+), 1312 deletions(-) create mode 100644 src/sparsevllm/distributed/sharding.py create mode 100644 src/sparsevllm/distributed/topology.py create mode 100644 src/sparsevllm/models/checkpoint.py create mode 100644 src/sparsevllm/models/layout.py create mode 100644 src/sparsevllm/models/spec.py create mode 100644 src/sparsevllm/quantization/config.py create mode 100644 src/sparsevllm/utils/config.py 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 4b1b794c..71819c7e 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -1,639 +1,64 @@ -"""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, + canonical_model_type, + 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_family_outer_config(raw_config): + model_type = canonical_model_type(config_get(raw_config, "model_type", "")) + 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 a non-Qwen3.5-family 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, *, 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}." - ) - 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}." - ) - - -_QWEN35_MOE_FIXED_FIELDS = { - "vocab_size": 248320, - "hidden_size": 2048, - "num_hidden_layers": 40, - "num_attention_heads": 16, - "num_key_value_heads": 2, - "head_dim": 256, - "linear_num_key_heads": 16, - "linear_num_value_heads": 32, - "linear_key_head_dim": 128, - "linear_value_head_dim": 128, - "linear_conv_kernel_dim": 4, - "num_experts": 256, - "num_experts_per_tok": 8, - "moe_intermediate_size": 512, - "shared_expert_intermediate_size": 512, - "max_position_embeddings": 262144, -} - - -def _validate_qwen35_moe_checkpoint_config( - outer_hf_config: Any, - hf_config: Any, - quantization_config: QuantizationConfig, -) -> None: - architectures = tuple( - _config_get(outer_hf_config, "architectures", ()) or () - ) - if architectures != ("Qwen3_5MoeForConditionalGeneration",): - raise ValueError( - "Qwen3.6 MoE requires " - "architectures=['Qwen3_5MoeForConditionalGeneration'], " - 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 NotImplementedError( - "Qwen3.6 MoE requires BF16 activations with either BF16 or block-FP8 " - "language-model weights, " - f"got dtype={configured_dtype!r}." - ) - for field_name, expected in _QWEN35_MOE_FIXED_FIELDS.items(): - actual = _config_get(hf_config, field_name, None) - if actual != expected: - raise ValueError( - f"Qwen3.6 MoE requires {field_name}={expected!r}, " - f"got {actual!r}." - ) - expected_values = { - "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, - } - for field_name, expected in expected_values.items(): - actual = _config_get(hf_config, field_name, None) - if actual != expected: - raise ValueError( - f"Qwen3.6 MoE requires {field_name}={expected!r}, " - f"got {actual!r}." - ) - layer_types = tuple(_config_get(hf_config, "layer_types", ()) or ()) - expected_layer_types = tuple( - "full_attention" if (layer_idx + 1) % 4 == 0 else "linear_attention" - for layer_idx in range(40) - ) - if layer_types != expected_layer_types: - raise ValueError( - "Qwen3.6 MoE requires the checkpoint's 3:1 Gated DeltaNet/full-" - "attention layer layout." - ) - - -_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 _is_qwen35_moe_outer_config(config: Any) -> bool: - return str(_config_get(config, "model_type", "") or "").strip().lower() in { - "qwen3_5_moe", - "qwen3_6_moe", - } - - -def _is_qwen35_family_outer_config(config: Any) -> bool: - return _is_qwen35_outer_config(config) or _is_qwen35_moe_outer_config(config) - - 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): @@ -641,299 +66,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_qwen35_moe_runtime(config, *, model_type: str) -> None: - outer_tp_size = int(config.tensor_parallel_size) - ep_size = int(config.expert_parallel_size) - if int(config.data_parallel_size) != 1: - raise ValueError( - "Qwen3.6 MoE requires DP=1, got " - f"TP={outer_tp_size}, EP={ep_size}, DP={config.data_parallel_size}." - ) - if outer_tp_size % ep_size: - raise ValueError( - "Qwen3.6 MoE outer tensor_parallel_size must be divisible by " - f"expert_parallel_size, got outer TP={outer_tp_size}, EP={ep_size}." - ) - hf_config = config.hf_config - num_experts = int(hf_config.num_experts) - if ep_size > num_experts or num_experts % ep_size: - raise ValueError( - "Qwen3.6 MoE num_experts must be divisible by expert_parallel_size " - f"and EP must not exceed experts, got experts={num_experts}, EP={ep_size}." - ) - attention_tp_fields = { - "num_attention_heads": int(hf_config.num_attention_heads), - "num_key_value_heads": int(hf_config.num_key_value_heads), - "linear_num_key_heads": int(hf_config.linear_num_key_heads), - "linear_num_value_heads": int(hf_config.linear_num_value_heads), - "vocab_size": int(hf_config.vocab_size), - "shared_expert_intermediate_size": int( - hf_config.shared_expert_intermediate_size - ), - } - invalid_attention_fields = { - name: value - for name, value in attention_tp_fields.items() - if value % outer_tp_size - } - if invalid_attention_fields: - raise ValueError( - "Qwen3.6 MoE attention/GDN/vocabulary/shared-expert dimensions " - "must be divisible by outer tensor_parallel_size, " - f"got TP={outer_tp_size}, invalid={invalid_attention_fields}." - ) - moe_tp_size = outer_tp_size // ep_size - if int(hf_config.moe_intermediate_size) % moe_tp_size: - raise ValueError( - "Qwen3.6 MoE moe_intermediate_size must be divisible by MoE TP " - f"size, got {hf_config.moe_intermediate_size} and {moe_tp_size}." - ) - top_k = int(hf_config.num_experts_per_tok) - if not 1 <= top_k <= num_experts: - raise ValueError( - "Qwen3.6 MoE num_experts_per_tok must be in [1, num_experts], " - f"got top_k={top_k}, num_experts={num_experts}." - ) - if getattr(hf_config, "torch_dtype", None) != torch.bfloat16: - raise NotImplementedError( - "Qwen3.6 MoE requires BF16 activations, got " - f"torch_dtype={getattr(hf_config, 'torch_dtype', None)}." - ) - if config.quantization_config.enabled: - block_size = tuple(config.quantization_config.weight_block_size or ()) - if block_size != (128, 128): - raise ValueError( - "Qwen3.6 MoE FP8 requires weight_block_size=(128, 128), " - f"got {block_size}." - ) - fp8_local_dimensions = { - "hidden_size": int(hf_config.hidden_size), - "shared_expert_intermediate_size": int( - hf_config.shared_expert_intermediate_size - ) - // outer_tp_size, - } - invalid_fp8_dimensions = { - name: value - for name, value in fp8_local_dimensions.items() - if value % 128 - } - if invalid_fp8_dimensions: - raise ValueError( - "Qwen3.6 MoE FP8 local Linear dimensions must be 128-aligned, " - f"got TP={outer_tp_size}, invalid={invalid_fp8_dimensions}." - ) - _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 @@ -942,47 +79,38 @@ 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) - is_qwen35_moe = _is_qwen35_moe_outer_config(config.outer_hf_config) - is_qwen35_family = is_qwen35 or is_qwen35_moe + config.outer_hf_config = _load_model_config(config.model) + model_type = canonical_model_type( + config_get(config.outer_hf_config, "model_type", "") + ) + 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") - elif is_qwen35_moe: - setattr(config.hf_config, "model_type", "qwen3_5_moe") - 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}." - ) - + setattr(config.hf_config, "model_type", model_type) + 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_family: + if not model_spec.supports_tiny_random: raise NotImplementedError( - "Tiny random mode does not support qwen3_5 family models yet." + f"Tiny random mode does not support {model_spec.name} yet." ) config.tiny_random_overrides = apply_tiny_random_overrides( config.hf_config, @@ -994,70 +122,38 @@ 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" - elif is_qwen35_moe: - quantized_model_name = "Qwen3.6 MoE" 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, + model_name=model_spec.name, ) - if is_qwen35: - _validate_qwen35_checkpoint_precision( - config.hf_config, - raw_quantization_config, - config.quantization_config, - ) - if is_qwen35_moe: - _validate_qwen35_moe_checkpoint_config( - config.outer_hf_config, - config.hf_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) + 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 == "qwen3_5_moe": - _validate_qwen35_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_family) - return is_qwen35_family + 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 38ebdb02..94b3895a 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.method_registry import OUTER_TP_MOE_MODEL_TYPES, PREFILL_POLICY_AUTO +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,11 @@ 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 OUTER_TP_MOE_MODEL_TYPES and int( - self.tensor_parallel_size - ) > 1 + return self.parallel_topology.is_outer_tp_moe @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 +99,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 +120,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/distributed/__init__.py b/src/sparsevllm/distributed/__init__.py index 9ce7f8c9..86588258 100644 --- a/src/sparsevllm/distributed/__init__.py +++ b/src/sparsevllm/distributed/__init__.py @@ -9,10 +9,14 @@ reset_parallel_context, world_rank_from_parallel_ranks, ) +from sparsevllm.distributed.topology import ParallelMode, ParallelTopology +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", @@ -20,4 +24,6 @@ "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 d418e4bf..9c1aaec4 100644 --- a/src/sparsevllm/distributed/parallel_context.py +++ b/src/sparsevllm/distributed/parallel_context.py @@ -5,6 +5,7 @@ import torch import torch.distributed as dist +from sparsevllm.distributed.topology import ParallelTopology from sparsevllm.operators.all_reduce import AllReduceProvider, resolve_all_reduce_provider @@ -117,22 +118,13 @@ def parallel_group_ranks( def hybrid_moe_group_ranks( *, - outer_tp_size: int, - moe_ep_size: int, + topology: ParallelTopology, ) -> 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 + if not topology.is_outer_tp_moe: + raise ValueError("Hybrid MoE groups require an Outer-TP MoE topology.") + outer_tp_size = topology.attention_tp_size + moe_ep_size = topology.expert_parallel_size + moe_tp_size = topology.moe_tp_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)) @@ -345,10 +337,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: @@ -356,10 +345,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: @@ -368,11 +357,8 @@ 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, - ) + if topology.is_outer_tp_moe: + hybrid_groups = hybrid_moe_group_ranks(topology=topology) ranks_by_dimension = { "tensor": hybrid_groups["attention"], "expert": hybrid_groups["moe_expert"], 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..676ada08 --- /dev/null +++ b/src/sparsevllm/distributed/topology.py @@ -0,0 +1,73 @@ +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" + + +@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(int(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}." + ) + + @property + def is_outer_tp_moe(self) -> bool: + return self.mode is ParallelMode.OUTER_TP_MOE + + @property + def attention_tp_size(self) -> int: + return self.tensor_parallel_size + + @property + def moe_tp_size(self) -> int: + 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: + return ( + self.tensor_parallel_size + if self.is_outer_tp_moe + else self.tensor_parallel_size + * self.expert_parallel_size + * self.data_parallel_size + ) diff --git a/src/sparsevllm/engine/llm_engine.py b/src/sparsevllm/engine/llm_engine.py index 2f7515d3..1ca2cff8 100644 --- a/src/sparsevllm/engine/llm_engine.py +++ b/src/sparsevllm/engine/llm_engine.py @@ -31,10 +31,8 @@ RequestAdmission, stable_token_digest, ) -from sparsevllm.method_registry import ( - OUTER_TP_MOE_MODEL_TYPES, - normalize_sparse_method, -) +from sparsevllm.method_registry import normalize_sparse_method +from sparsevllm.models.spec import resolve_model_spec from sparsevllm.utils.profiler import profiler def _deltakv_graph_warmup_profile(config: Config) -> str: @@ -66,8 +64,9 @@ 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 "") - has_experts = int(getattr(config.hf_config, "num_experts", 0) or 0) > 0 - if not has_experts and model_type not in OUTER_TP_MOE_MODEL_TYPES: + if not model_type: + return () + if resolve_model_spec(model_type).num_experts_field is None: return () max_batched_tokens = int(config.max_num_batched_tokens) diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index fd481639..cce0ac03 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -30,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 resolve_model_spec import sparsevllm.platforms as platforms from sparsevllm.utils.profiler import profiler @@ -45,24 +46,27 @@ 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 - _QWEN35_MOE_IMPORT_ERROR = None -except ImportError as exc: +except ImportError: Qwen35MoeForCausalLM = None - _QWEN35_MOE_IMPORT_ERROR = exc + + +def _create_model(hf_config): + model_spec = resolve_model_spec(hf_config.model_type) + 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_" @@ -140,10 +144,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(). @@ -161,49 +162,7 @@ def __init__( bool(getattr(config, "decode_cuda_graph", 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 == "qwen3_5_moe": - if Qwen35MoeForCausalLM is None: - raise ImportError( - "Qwen35MoeForCausalLM is unavailable; verify the Qwen3.6 MoE " - f"runtime imports: {_QWEN35_MOE_IMPORT_ERROR}" - ) from _QWEN35_MOE_IMPORT_ERROR - self.model = Qwen35MoeForCausalLM(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) if config.tiny_random: from sparsevllm.debug.tiny_random import initialize_sparse_model diff --git a/src/sparsevllm/method_registry.py b/src/sparsevllm/method_registry.py index 07e81a19..ab280cac 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,94 +59,56 @@ "skipkv", } -H2O_SUPPORTED_MODEL_TYPES = frozenset( +H2O_SUPPORTED_MODEL_TYPES = frozenset(MODEL_SPECS) + +SKIPKV_ASSET_MODEL_NAMES = frozenset( { - "qwen2", - "qwen3", - "qwen3_moe", - "qwen3_5", - "qwen3_5_moe", - "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 QWEN35_MOE_COMPATIBILITY = ModelRuntimeCompatibility( - parallel_mode="outer_tp_moe_tp_ep", sparse_methods=QWEN3_MOE_TP_EP_COMPATIBILITY.sparse_methods, prefix_cache_methods=frozenset({""}), requires_eager=False, @@ -153,37 +118,13 @@ class ModelRuntimeCompatibility: ) MINIMAX_M2_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"}), 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, @@ -191,19 +132,16 @@ class ModelRuntimeCompatibility: ) MODEL_RUNTIME_COMPATIBILITY = { - "qwen3_moe": QWEN3_MOE_EP_COMPATIBILITY, - "qwen3_5_moe": QWEN35_MOE_COMPATIBILITY, - "minimax_m2": MINIMAX_M2_EP_COMPATIBILITY, -} - -OUTER_TP_MOE_MODEL_TYPES = frozenset( - {"qwen3_moe", "qwen3_5_moe", "minimax_m2"} -) - -OUTER_TP_RUNTIME_COMPATIBILITY = { - "qwen3_moe": QWEN3_MOE_TP_EP_COMPATIBILITY, - "qwen3_5_moe": QWEN35_MOE_COMPATIBILITY, - "minimax_m2": MINIMAX_M2_TP_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, } # All shipped cache managers now expose a graph-stable decode preparation path. @@ -256,6 +194,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" @@ -272,38 +222,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 OUTER_TP_MOE_MODEL_TYPES and tp_size > 1: - compatibility = OUTER_TP_RUNTIME_COMPATIBILITY[model_type] - 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: @@ -315,22 +247,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..89ab4cf3 --- /dev/null +++ b/src/sparsevllm/models/checkpoint.py @@ -0,0 +1,279 @@ +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", + ) + + +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) + + +CHECKPOINT_VALIDATORS = { + "qwen3": _qwen3_checkpoint, + "qwen3_moe": _qwen3_moe_checkpoint, + "qwen3_5": _qwen35_checkpoint, + "qwen3_5_moe": _qwen35_moe_checkpoint, + "minimax_m2": _minimax_checkpoint, +} diff --git a/src/sparsevllm/models/layout.py b/src/sparsevllm/models/layout.py new file mode 100644 index 00000000..f1b7eb9d --- /dev/null +++ b/src/sparsevllm/models/layout.py @@ -0,0 +1,219 @@ +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", + }: + 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/spec.py b/src/sparsevllm/models/spec.py new file mode 100644 index 00000000..a4427f9e --- /dev/null +++ b/src/sparsevllm/models/spec.py @@ -0,0 +1,184 @@ +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 + prefix_cache_block_size_multiple: int | None = None + deltakv_checkpoint_model_types: frozenset[str] = frozenset() + runtime_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: + topology = ParallelTopology( + int(tp_size), + int(ep_size), + int(dp_size), + ( + ParallelMode.OUTER_TP_MOE + if self.supports_outer_tp_moe and int(tp_size) > 1 + else ParallelMode.STANDARD + ), + ) + 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", "qwen3_6"}), + 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", "qwen3_6"}), + 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", + ), + } +) + + +MODEL_TYPE_ALIASES = { + "qwen3_6": "qwen3_5", + "qwen3_6_moe": "qwen3_5_moe", +} + + +def canonical_model_type(model_type: str | None) -> str: + normalized = str(model_type or "").strip().lower() + return MODEL_TYPE_ALIASES.get(normalized, normalized) + + +def resolve_model_spec(model_type: str) -> ModelSpec: + model_type = canonical_model_type(model_type) + 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/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/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/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_parallel_context.py b/tests/test_parallel_context.py index a914e7b5..8e2c9490 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -6,8 +6,13 @@ import torch.distributed as dist import sparsevllm.platforms as platforms -from sparsevllm.config import Config -from sparsevllm.distributed import ParallelContext, ParallelGroup +from sparsevllm.config import Config, RuntimeLayout +from sparsevllm.distributed import ( + ParallelContext, + ParallelGroup, + ParallelMode, + ParallelTopology, +) from sparsevllm.distributed.parallel_context import ( get_parallel_context, hybrid_moe_group_ranks, @@ -112,7 +117,9 @@ def test_parallel_group_members_follow_dp_ep_tp_layout(): def test_hybrid_moe_groups_split_outer_attention_world(): - assert hybrid_moe_group_ranks(outer_tp_size=4, moe_ep_size=2) == { + assert hybrid_moe_group_ranks( + topology=ParallelTopology(4, 2, 1, ParallelMode.OUTER_TP_MOE) + ) == { "attention": ((0, 1, 2, 3),), "moe_tensor": ((0, 1), (2, 3)), "moe_expert": ((0, 2), (1, 3)), @@ -120,6 +127,27 @@ def test_hybrid_moe_groups_split_outer_attention_world(): } +def test_parallel_topology_resolves_rank_local_sizes(): + standard = ParallelTopology(2, 4, 1) + hybrid = ParallelTopology(4, 2, 1, ParallelMode.OUTER_TP_MOE) + + 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) + + +@pytest.mark.parametrize( + "topology", + [ + (0, 1, 1, ParallelMode.STANDARD), + (4, 3, 1, ParallelMode.OUTER_TP_MOE), + (4, 2, 2, ParallelMode.OUTER_TP_MOE), + ], +) +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 ( @@ -130,10 +158,7 @@ def test_hybrid_moe_parallel_context_uses_explicit_groups(): 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 @@ -160,7 +185,8 @@ def new_group(ranks): 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 @@ -170,7 +196,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), @@ -191,7 +217,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(): @@ -237,7 +263,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() @@ -247,7 +273,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)): @@ -354,7 +380,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) @@ -401,7 +427,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 9fba25b8..24a5d22b 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") diff --git a/tests/test_qwen35_mixed_runtime.py b/tests/test_qwen35_mixed_runtime.py index b168f81c..d9f660b0 100644 --- a/tests/test_qwen35_mixed_runtime.py +++ b/tests/test_qwen35_mixed_runtime.py @@ -8,8 +8,8 @@ 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, @@ -46,6 +46,7 @@ _get_rotary_dim, ) from sparsevllm.models.qwen3_5_moe import Qwen35MoeRouter, Qwen35MoeSparseMoeBlock +from sparsevllm.models.checkpoint import validate_checkpoint 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 @@ -98,6 +99,36 @@ 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 @@ -872,6 +903,7 @@ 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), @@ -1167,7 +1199,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", @@ -1215,7 +1247,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)) @@ -1224,7 +1256,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 5b8bd379..dd3f88e5 100644 --- a/tests/test_qwen3_moe_compatibility.py +++ b/tests/test_qwen3_moe_compatibility.py @@ -1,6 +1,8 @@ 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, @@ -22,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", @@ -64,9 +77,7 @@ def test_qwen35_moe_registry_accepts_vanilla_prefix_cache(): assert validate_model_runtime_compatibility( model_type="qwen3_5_moe", sparse_method="", - tensor_parallel_size=2, - expert_parallel_size=2, - data_parallel_size=1, + topology=ParallelTopology(2, 2, 1, ParallelMode.OUTER_TP_MOE), enforce_eager=True, decode_cuda_graph=False, enable_prefix_caching=True, @@ -93,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") @@ -154,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"} From bec54eca4161c2f487c222c8a1bdde1ddf394afb Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 14:24:09 +0800 Subject: [PATCH 30/35] refactor: require validated model layout --- src/sparsevllm/engine/cache_manager/base.py | 8 +------- src/sparsevllm/utils/select_omnikv_full_layers.py | 15 ++++----------- tests/test_snapkv_cache_budget.py | 3 ++- tests/test_sparsevllm_regression_grading.py | 2 ++ 4 files changed, 9 insertions(+), 19 deletions(-) diff --git a/src/sparsevllm/engine/cache_manager/base.py b/src/sparsevllm/engine/cache_manager/base.py index 035eb315..d4fffd6c 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,10 +266,6 @@ 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}.") if sparse_method == "deltakv": 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_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 7e5a1d41..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, From 01ed8f9ed642d98dc24b696c8916a744fff02a0d Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 14:34:02 +0800 Subject: [PATCH 31/35] refactor: separate parallel topology --- src/sparsevllm/distributed/__init__.py | 9 +- .../distributed/parallel_context.py | 152 +----------------- src/sparsevllm/distributed/topology.py | 105 +++++++++++- tests/test_parallel_context.py | 33 ++-- 4 files changed, 122 insertions(+), 177 deletions(-) diff --git a/src/sparsevllm/distributed/__init__.py b/src/sparsevllm/distributed/__init__.py index 86588258..3f8700fc 100644 --- a/src/sparsevllm/distributed/__init__.py +++ b/src/sparsevllm/distributed/__init__.py @@ -2,14 +2,16 @@ 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.topology import ParallelMode, ParallelTopology from sparsevllm.distributed.sharding import validate_model_sharding, validate_top_k __all__ = [ @@ -18,7 +20,6 @@ "ParallelMode", "ParallelTopology", "get_parallel_context", - "hybrid_moe_group_ranks", "init_parallel_context", "parallel_group_ranks", "parallel_ranks_from_world_rank", diff --git a/src/sparsevllm/distributed/parallel_context.py b/src/sparsevllm/distributed/parallel_context.py index 9c1aaec4..eb41ae37 100644 --- a/src/sparsevllm/distributed/parallel_context.py +++ b/src/sparsevllm/distributed/parallel_context.py @@ -5,144 +5,10 @@ import torch import torch.distributed as dist -from sparsevllm.distributed.topology import ParallelTopology +from sparsevllm.distributed.topology import ParallelTopology, parallel_group_ranks from sparsevllm.operators.all_reduce import AllReduceProvider, resolve_all_reduce_provider -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( - *, - topology: ParallelTopology, -) -> dict[str, tuple[tuple[int, ...], ...]]: - if not topology.is_outer_tp_moe: - raise ValueError("Hybrid MoE groups require an Outer-TP MoE topology.") - outer_tp_size = topology.attention_tp_size - moe_ep_size = topology.expert_parallel_size - moe_tp_size = topology.moe_tp_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, - } - - @dataclass(frozen=True) class ParallelGroup: process_group: dist.ProcessGroup | None @@ -357,21 +223,7 @@ def init_parallel_context( f"world_size={world_size}, TP={tp_size}, EP={ep_size}, DP={dp_size}." ) - if topology.is_outer_tp_moe: - hybrid_groups = hybrid_moe_group_ranks(topology=topology) - 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, diff --git a/src/sparsevllm/distributed/topology.py b/src/sparsevllm/distributed/topology.py index 676ada08..e74096e5 100644 --- a/src/sparsevllm/distributed/topology.py +++ b/src/sparsevllm/distributed/topology.py @@ -29,7 +29,7 @@ def __post_init__(self) -> None: self.expert_parallel_size, self.data_parallel_size, ) - if any(int(size) <= 0 for size in sizes): + 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]}." @@ -71,3 +71,106 @@ def world_size(self) -> int: * 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.is_outer_tp_moe: + raise ValueError("Outer-TP MoE 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.is_outer_tp_moe: + raise ValueError("Outer-TP MoE 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 parallel_group_ranks( + topology: ParallelTopology, +) -> dict[str, tuple[tuple[int, ...], ...]]: + return ( + _outer_tp_moe_group_ranks(topology) + if topology.is_outer_tp_moe + else _standard_group_ranks(topology) + ) diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py index 8e2c9490..e46c67a3 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -12,10 +12,7 @@ ParallelGroup, ParallelMode, ParallelTopology, -) -from sparsevllm.distributed.parallel_context import ( get_parallel_context, - hybrid_moe_group_ranks, init_parallel_context, parallel_group_ranks, parallel_ranks_from_world_rank, @@ -93,36 +90,28 @@ 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( - topology=ParallelTopology(4, 2, 1, ParallelMode.OUTER_TP_MOE) - ) == { - "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,)), } From fa480e9d13d7a1fc1b094a76dd7fefd46d18648d Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 14:45:00 +0800 Subject: [PATCH 32/35] test: consolidate operator coverage --- tests/test_all_reduce_operator.py | 17 ----- tests/test_gate_up_swiglu_kernel.py | 28 -------- tests/test_gate_up_swiglu_operator.py | 85 ------------------------ tests/test_operator_providers.py | 93 +++++++++++++++++++++++++-- tests/test_qwen35_gated_shared_add.py | 21 ------ tests/test_triton_moe.py | 41 ++++++++++++ tests/test_vllm_microbench.py | 32 --------- 7 files changed, 129 insertions(+), 188 deletions(-) delete mode 100644 tests/test_all_reduce_operator.py delete mode 100644 tests/test_gate_up_swiglu_kernel.py delete mode 100644 tests/test_gate_up_swiglu_operator.py delete mode 100644 tests/test_qwen35_gated_shared_add.py delete mode 100644 tests/test_vllm_microbench.py diff --git a/tests/test_all_reduce_operator.py b/tests/test_all_reduce_operator.py deleted file mode 100644 index b46fb5e8..00000000 --- a/tests/test_all_reduce_operator.py +++ /dev/null @@ -1,17 +0,0 @@ -from unittest.mock import Mock - -import torch - -from sparsevllm.operators.all_reduce import HopperTp2FlashInferAllReduceProvider - - -def test_flashinfer_all_reduce_dispatches_unsupported_shape_before_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) diff --git a/tests/test_gate_up_swiglu_kernel.py b/tests/test_gate_up_swiglu_kernel.py deleted file mode 100644 index e9e0f926..00000000 --- a/tests/test_gate_up_swiglu_kernel.py +++ /dev/null @@ -1,28 +0,0 @@ -import pytest -import torch - -from sparsevllm.triton_kernel.gate_up_swiglu import h20_gate_up_swiglu - - -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) diff --git a/tests/test_gate_up_swiglu_operator.py b/tests/test_gate_up_swiglu_operator.py deleted file mode 100644 index 3d0cef5d..00000000 --- a/tests/test_gate_up_swiglu_operator.py +++ /dev/null @@ -1,85 +0,0 @@ -import pytest -import torch - -from sparsevllm.operators.gate_up_swiglu import ( - GATE_UP_SWIGLU_REGISTRY, - GateUpSwiGLUOpSpec, - NativeGateUpSwiGLUProvider, -) -from sparsevllm.operators.registry import OpResolver -from sparsevllm.platforms import DeviceCaps, PlatformEnum - - -def _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 _caps(device_name="NVIDIA H20", capability=(9, 0)) -> DeviceCaps: - return DeviceCaps( - platform=PlatformEnum.CUDA, - device_type="cuda", - device_index=0, - device_name=device_name, - compute_capability=capability, - runtime_version="13.0", - supports_graph_capture=True, - supports_triton=True, - supports_bfloat16=True, - supports_native_fp8=True, - ) - - -@pytest.mark.parametrize("tp_size", [1, 2]) -def test_h20_provider_requires_profiled_qwen36_shape(tp_size): - resolved = OpResolver(GATE_UP_SWIGLU_REGISTRY).resolve( - _spec(tp_size=tp_size), _caps() - ) - - assert resolved.provider.name == "h20_triton_decode" - - -@pytest.mark.parametrize( - "spec,caps", - [ - (_spec(weight_dtype=torch.float8_e4m3fn), _caps()), - (_spec(intermediate_size=768), _caps()), - (_spec(), _caps(device_name="NVIDIA H100 80GB HBM3")), - (_spec(), _caps(capability=(8, 9))), - ], -) -def test_unprofiled_shape_uses_native_provider(spec, caps): - resolved = OpResolver(GATE_UP_SWIGLU_REGISTRY).resolve(spec, caps) - - assert resolved.provider.name == "native" - - -def test_native_provider_matches_gate_up_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( - _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) diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 9d34ab44..f6357d45 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,10 +13,14 @@ TritonFp8LinearProvider, resolve_fp8_linear_provider, ) +from sparsevllm.operators.gate_up_swiglu import ( + GATE_UP_SWIGLU_REGISTRY, + GateUpSwiGLUOpSpec, + NativeGateUpSwiGLUProvider, +) from sparsevllm.operators.moe import ( MOE_REGISTRY, FlashInferCutlassFp8MoeProvider, - H20Qwen36HybridFp8MoeProvider, HopperQwen36HybridFp8MoeProvider, MoeOpSpec, resolve_moe_provider, @@ -105,6 +110,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", [ @@ -687,10 +774,6 @@ def test_h20_qwen36_hybrid_moe_uses_profiled_provider(): assert resolved.provider.name == "h20_qwen36_hybrid_fp8" -def test_h20_qwen36_hybrid_moe_limits_triton_to_profiled_token_count(): - assert H20Qwen36HybridFp8MoeProvider.TRITON_MAX_TOKENS_BY_EP_SIZE == {1: 8, 2: 1} - - def test_qwen36_hybrid_moe_dispatches_by_token_bucket(): provider = HopperQwen36HybridFp8MoeProvider() spec = _moe_spec( diff --git a/tests/test_qwen35_gated_shared_add.py b/tests/test_qwen35_gated_shared_add.py deleted file mode 100644 index 71dda157..00000000 --- a/tests/test_qwen35_gated_shared_add.py +++ /dev/null @@ -1,21 +0,0 @@ -import pytest -import torch - -from sparsevllm.operators.gated_shared_add import gated_shared_add - - -pytestmark = 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) diff --git a/tests/test_triton_moe.py b/tests/test_triton_moe.py index 1ab64422..a0c29f5a 100644 --- a/tests/test_triton_moe.py +++ b/tests/test_triton_moe.py @@ -4,6 +4,8 @@ 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, @@ -13,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, diff --git a/tests/test_vllm_microbench.py b/tests/test_vllm_microbench.py deleted file mode 100644 index 41d20985..00000000 --- a/tests/test_vllm_microbench.py +++ /dev/null @@ -1,32 +0,0 @@ -import argparse -from types import SimpleNamespace - -import pytest - -from benchmark.vllm_microbench import _parse_positive_ints, _validate_args - - -def test_vllm_microbench_parses_unique_batch_sizes(): - assert _parse_positive_ints("1,2,4") == [1, 2, 4] - - -@pytest.mark.parametrize("value", ["", "0,1", "1,1"]) -def test_vllm_microbench_rejects_invalid_batch_sizes(value): - with pytest.raises(argparse.ArgumentTypeError): - _parse_positive_ints(value) - - -def test_vllm_microbench_rejects_short_model_context(): - args = SimpleNamespace( - input_len=1024, - output_len=128, - num_warmups=2, - num_iters=5, - tensor_parallel_size=2, - max_model_len=1151, - max_num_batched_tokens=4096, - gpu_memory_utilization=0.7, - ) - - with pytest.raises(ValueError, match=r"input_len \+ output_len"): - _validate_args(args) From 8adb8c6abca7d0e15e9618f0d3b0e4b984332f1a Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 14:46:13 +0800 Subject: [PATCH 33/35] docs: clarify optional qwen dependencies --- README.md | 2 -- README_zh.md | 13 ++++++++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 5b28e54a..39f8b2f0 100644 --- a/README.md +++ b/README.md @@ -159,8 +159,6 @@ the smaller CUDA-specific extra: pip install -e ".[prefix-offload]" ``` -Sparse-vLLM supports Qwen3.5/Qwen3.6 dense and MoE 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 b8351cfc..021af344 100644 --- a/README_zh.md +++ b/README_zh.md @@ -128,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 @@ -138,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)。 ## 基准测试 From 187c69b88c9ad6b639b58007d96757049019cd20 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 14:53:11 +0800 Subject: [PATCH 34/35] refactor: preserve checkpoint model types --- src/sparsevllm/configs/model.py | 13 +++---------- src/sparsevllm/engine/llm_engine.py | 6 +----- src/sparsevllm/engine/model_runner.py | 7 +++---- src/sparsevllm/models/spec.py | 16 ++-------------- tests/test_prefill_schedule_policy.py | 9 +++++---- tests/test_qwen35_mixed_runtime.py | 4 +++- 6 files changed, 17 insertions(+), 38 deletions(-) diff --git a/src/sparsevllm/configs/model.py b/src/sparsevllm/configs/model.py index 71819c7e..cf4a598d 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -13,11 +13,7 @@ ) from sparsevllm.models.checkpoint import validate_checkpoint from sparsevllm.models.layout import RuntimeLayout -from sparsevllm.models.spec import ( - ModelSpec, - canonical_model_type, - resolve_model_spec, -) +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 @@ -41,7 +37,7 @@ def _load_model_config(model_path: str) -> Any: ) from load_error with open(config_path, "r", encoding="utf-8") as f: raw_config = json.load(f) - model_type = canonical_model_type(config_get(raw_config, "model_type", "")) + 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( @@ -93,12 +89,9 @@ def load_and_validate_model(config) -> None: "Tiny random mode does not support DeltaKV compressor weights yet." ) config.outer_hf_config = _load_model_config(config.model) - model_type = canonical_model_type( - config_get(config.outer_hf_config, "model_type", "") - ) + 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) - setattr(config.hf_config, "model_type", model_type) config.model_spec = model_spec config.parallel_topology = model_spec.topology( config.tensor_parallel_size, diff --git a/src/sparsevllm/engine/llm_engine.py b/src/sparsevllm/engine/llm_engine.py index 1ca2cff8..a18b3465 100644 --- a/src/sparsevllm/engine/llm_engine.py +++ b/src/sparsevllm/engine/llm_engine.py @@ -32,7 +32,6 @@ stable_token_digest, ) from sparsevllm.method_registry import normalize_sparse_method -from sparsevllm.models.spec import resolve_model_spec from sparsevllm.utils.profiler import profiler def _deltakv_graph_warmup_profile(config: Config) -> str: @@ -63,10 +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 not model_type: - return () - if resolve_model_spec(model_type).num_experts_field is None: + if config.model_spec.num_experts_field is None: return () max_batched_tokens = int(config.max_num_batched_tokens) diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index cce0ac03..7caf2f42 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -30,7 +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 resolve_model_spec +from sparsevllm.models.spec import ModelSpec import sparsevllm.platforms as platforms from sparsevllm.utils.profiler import profiler @@ -60,8 +60,7 @@ Qwen35MoeForCausalLM = None -def _create_model(hf_config): - model_spec = resolve_model_spec(hf_config.model_type) +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: @@ -162,7 +161,7 @@ def __init__( bool(getattr(config, "decode_cuda_graph", False)), ) - self.model = _create_model(hf_config) + self.model = _create_model(hf_config, config.model_spec) if config.tiny_random: from sparsevllm.debug.tiny_random import initialize_sparse_model diff --git a/src/sparsevllm/models/spec.py b/src/sparsevllm/models/spec.py index a4427f9e..439ef192 100644 --- a/src/sparsevllm/models/spec.py +++ b/src/sparsevllm/models/spec.py @@ -116,7 +116,7 @@ def validate_sharding(self, hf_config: Any, topology: ParallelTopology) -> None: allow_raw_config=True, supports_tiny_random=False, prefix_cache_block_size_multiple=4096, - deltakv_checkpoint_model_types=frozenset({"qwen3_5", "qwen3_6"}), + deltakv_checkpoint_model_types=frozenset({"qwen3_5"}), runtime_class_name="Qwen35ForCausalLM", attention_tp_fields=_QWEN35_TP_FIELDS, ), @@ -138,7 +138,7 @@ def validate_sharding(self, hf_config: Any, topology: ParallelTopology) -> None: supports_expert_parallel=True, supports_outer_tp_moe=True, prefix_cache_block_size_multiple=4096, - deltakv_checkpoint_model_types=frozenset({"qwen3_5", "qwen3_6"}), + deltakv_checkpoint_model_types=frozenset({"qwen3_5"}), runtime_class_name="Qwen35MoeForCausalLM", attention_tp_fields=( *_QWEN35_TP_FIELDS, @@ -163,19 +163,7 @@ def validate_sharding(self, hf_config: Any, topology: ParallelTopology) -> None: ) -MODEL_TYPE_ALIASES = { - "qwen3_6": "qwen3_5", - "qwen3_6_moe": "qwen3_5_moe", -} - - -def canonical_model_type(model_type: str | None) -> str: - normalized = str(model_type or "").strip().lower() - return MODEL_TYPE_ALIASES.get(normalized, normalized) - - def resolve_model_spec(model_type: str) -> ModelSpec: - model_type = canonical_model_type(model_type) if model_type not in MODEL_SPECS: supported = ", ".join(sorted(MODEL_SPECS)) raise NotImplementedError( diff --git a/tests/test_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index 24a5d22b..969359a6 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -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 @@ -1686,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), @@ -1695,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), ()) @@ -1705,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( @@ -1728,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 d9f660b0..75a54a29 100644 --- a/tests/test_qwen35_mixed_runtime.py +++ b/tests/test_qwen35_mixed_runtime.py @@ -47,6 +47,7 @@ ) 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 @@ -907,6 +908,7 @@ def stop_at_model_construction(_config): 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), @@ -1225,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 From 02ec8681d3867743402d5afd2767edd06124441a Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 14:57:24 +0800 Subject: [PATCH 35/35] feat: support deepseek v4 flash --- .../debug/deepseek_v4_flash_tiny_random.json | 18 + docs/zh/benchmarking/README.md | 1 + docs/zh/benchmarking/deepseek-v4-flash.md | 110 +++ docs/zh/design/README.md | 1 + docs/zh/design/deepseek-v4-flash.md | 92 +++ docs/zh/features/supported-models.md | 2 + src/sparsevllm/configs/model.py | 26 +- src/sparsevllm/configs/runtime.py | 4 + src/sparsevllm/debug/tiny_random.py | 32 +- .../distributed/parallel_context.py | 35 + src/sparsevllm/distributed/topology.py | 55 +- .../engine/cache_manager/__init__.py | 5 + src/sparsevllm/engine/cache_manager/base.py | 16 + .../engine/cache_manager/deepseek_v4.py | 170 +++++ src/sparsevllm/engine/model_runner.py | 115 ++- src/sparsevllm/method_registry.py | 16 +- src/sparsevllm/models/checkpoint.py | 69 ++ src/sparsevllm/models/deepseek_v4.py | 293 ++++++++ src/sparsevllm/models/deepseek_v4_native.py | 676 ++++++++++++++++++ src/sparsevllm/models/layout.py | 2 + src/sparsevllm/models/spec.py | 30 +- src/sparsevllm/operators/moe.py | 156 ++++ src/sparsevllm/utils/loader.py | 69 +- tests/test_deepseek_v4_config.py | 162 +++++ tests/test_deepseek_v4_model.py | 359 ++++++++++ tests/test_operator_providers.py | 57 ++ tests/test_parallel_context.py | 43 +- tests/test_weight_loading.py | 48 ++ 28 files changed, 2599 insertions(+), 63 deletions(-) create mode 100644 configs/debug/deepseek_v4_flash_tiny_random.json create mode 100644 docs/zh/benchmarking/deepseek-v4-flash.md create mode 100644 docs/zh/design/deepseek-v4-flash.md create mode 100644 src/sparsevllm/engine/cache_manager/deepseek_v4.py create mode 100644 src/sparsevllm/models/deepseek_v4.py create mode 100644 src/sparsevllm/models/deepseek_v4_native.py create mode 100644 tests/test_deepseek_v4_config.py create mode 100644 tests/test_deepseek_v4_model.py 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/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 44331922..9b67d2ee 100644 --- a/docs/zh/features/supported-models.md +++ b/docs/zh/features/supported-models.md @@ -15,6 +15,7 @@ | 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`。 @@ -40,6 +41,7 @@ MoE 使用 `model_type=qwen3_5_moe`。 | 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/model.py b/src/sparsevllm/configs/model.py index cf4a598d..e808ac7e 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -122,9 +122,14 @@ def load_and_validate_model(config) -> None: "quantization_config", config_get(config.outer_hf_config, "quantization_config", None), ) + 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=model_spec.requires_fp8, + required_fp8=model_spec.requires_fp8 and not use_unquantized_tiny, model_name=model_spec.name, ) if config.tiny_random and config.quantization_config.enabled: @@ -132,14 +137,17 @@ def load_and_validate_model(config) -> None: "Tiny random mode does not support quantized model weights." ) setattr(config.hf_config, "quantization_config", config.quantization_config) - 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 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, + ) validate_model_runtime_compatibility( model_type=model_type, diff --git a/src/sparsevllm/configs/runtime.py b/src/sparsevllm/configs/runtime.py index 94b3895a..25459d59 100644 --- a/src/sparsevllm/configs/runtime.py +++ b/src/sparsevllm/configs/runtime.py @@ -89,6 +89,10 @@ class Config( def uses_outer_tp_moe_layout(self) -> bool: 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 self.parallel_topology.attention_tp_size 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/parallel_context.py b/src/sparsevllm/distributed/parallel_context.py index eb41ae37..cee2da05 100644 --- a/src/sparsevllm/distributed/parallel_context.py +++ b/src/sparsevllm/distributed/parallel_context.py @@ -129,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, diff --git a/src/sparsevllm/distributed/topology.py b/src/sparsevllm/distributed/topology.py index e74096e5..327f6e9e 100644 --- a/src/sparsevllm/distributed/topology.py +++ b/src/sparsevllm/distributed/topology.py @@ -7,6 +7,7 @@ class ParallelMode(str, Enum): STANDARD = "standard" OUTER_TP_MOE = "outer_tp_moe_tp_ep" + DPA_EP = "dpa_ep" @dataclass(frozen=True) @@ -45,17 +46,32 @@ def __post_init__(self) -> None: "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 self.tensor_parallel_size + 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 @@ -64,6 +80,8 @@ def moe_tp_size(self) -> int: @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 @@ -79,8 +97,10 @@ def world_rank_from_parallel_ranks( ep_rank: int, tp_rank: int, ) -> int: - if topology.is_outer_tp_moe: - raise ValueError("Outer-TP MoE does not use standard DP/EP/TP rank mapping.") + 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), @@ -100,8 +120,10 @@ def parallel_ranks_from_world_rank( topology: ParallelTopology, world_rank: int, ) -> tuple[int, int, int]: - if topology.is_outer_tp_moe: - raise ValueError("Outer-TP MoE does not use standard DP/EP/TP rank mapping.") + 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( @@ -166,11 +188,24 @@ def _outer_tp_moe_group_ranks( } +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, ...], ...]]: - return ( - _outer_tp_moe_group_ranks(topology) - if topology.is_outer_tp_moe - else _standard_group_ranks(topology) - ) + 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 d4fffd6c..cc8bc6a7 100644 --- a/src/sparsevllm/engine/cache_manager/base.py +++ b/src/sparsevllm/engine/cache_manager/base.py @@ -268,6 +268,22 @@ def create(config: Config, parallel_context: ParallelContext) -> "CacheManager": sparse_method = normalize_sparse_method(config.vllm_sparse_method) 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/model_runner.py b/src/sparsevllm/engine/model_runner.py index 7caf2f42..bd0eb0b1 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -59,6 +59,11 @@ 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 @@ -160,6 +165,11 @@ def __init__( "decode_cuda_graph", bool(getattr(config, "decode_cuda_graph", False)), ) + setattr( + hf_config, + "sparsevllm_tiny_random", + bool(getattr(config, "tiny_random", False)), + ) self.model = _create_model(hf_config, config.model_spec) if config.tiny_random: @@ -1022,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: @@ -1038,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 ( @@ -1233,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/method_registry.py b/src/sparsevllm/method_registry.py index ab280cac..e0758886 100644 --- a/src/sparsevllm/method_registry.py +++ b/src/sparsevllm/method_registry.py @@ -59,8 +59,6 @@ "skipkv", } -H2O_SUPPORTED_MODEL_TYPES = frozenset(MODEL_SPECS) - SKIPKV_ASSET_MODEL_NAMES = frozenset( { "DeepSeek-R1-Distill-Llama-8B", @@ -117,6 +115,13 @@ class ModelRuntimeCompatibility: ), ) +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"}), @@ -142,8 +147,15 @@ class ModelRuntimeCompatibility: ("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 = { diff --git a/src/sparsevllm/models/checkpoint.py b/src/sparsevllm/models/checkpoint.py index 89ab4cf3..a8d545e4 100644 --- a/src/sparsevllm/models/checkpoint.py +++ b/src/sparsevllm/models/checkpoint.py @@ -228,6 +228,70 @@ def _validate_minimax(config: Any, raw_quantization_config: Any) -> None: ) +_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, *, @@ -270,10 +334,15 @@ 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 index f1b7eb9d..91890743 100644 --- a/src/sparsevllm/models/layout.py +++ b/src/sparsevllm/models/layout.py @@ -34,6 +34,8 @@ def _attention_type(value: Any) -> str: "attention", "self_attention", "sliding_attention", + "compressed_sparse_attention", + "heavily_compressed_attention", }: return "full" if value in { diff --git a/src/sparsevllm/models/spec.py b/src/sparsevllm/models/spec.py index 439ef192..c7547acb 100644 --- a/src/sparsevllm/models/spec.py +++ b/src/sparsevllm/models/spec.py @@ -18,24 +18,29 @@ class ModelSpec: 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), - ( - ParallelMode.OUTER_TP_MOE - if self.supports_outer_tp_moe and int(tp_size) > 1 - else ParallelMode.STANDARD - ), + mode, ) if topology.expert_parallel_size > 1 and not self.supports_expert_parallel: raise ValueError( @@ -159,6 +164,21 @@ def validate_sharding(self, hf_config: Any, topology: ParallelTopology) -> None: 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", + ), } ) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 5ab1ae7e..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,6 +352,161 @@ 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" 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/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_operator_providers.py b/tests/test_operator_providers.py index f6357d45..2805fd0c 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -20,6 +20,7 @@ ) from sparsevllm.operators.moe import ( MOE_REGISTRY, + FlashInferCutlassFp4MoeProvider, FlashInferCutlassFp8MoeProvider, HopperQwen36HybridFp8MoeProvider, MoeOpSpec, @@ -77,6 +78,7 @@ def _moe_spec( routing_method="softmax", scale_dtype=None, cuda_graph=True, + activation_limit=None, ) -> MoeOpSpec: return MoeOpSpec( num_experts=num_experts, @@ -92,6 +94,7 @@ def _moe_spec( tp_size=tp_size, routing_method=routing_method, scale_dtype=scale_dtype, + activation_limit=activation_limit, ) @@ -639,6 +642,60 @@ 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, diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py index e46c67a3..e2d91b2c 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -116,12 +116,24 @@ def test_hybrid_moe_groups_split_outer_attention_world(): } +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( @@ -130,6 +142,8 @@ def test_parallel_topology_resolves_rank_local_sizes(): (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): @@ -141,10 +155,10 @@ 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()), + 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(4, 2, 1, ParallelMode.OUTER_TP_MOE), @@ -158,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 = [] 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)