diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index c7ef61b719..1bfe200ade 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -416,7 +416,7 @@ steps: - tests/kernels/helion/ - vllm/platforms/rocm.py commands: - - pip install helion==1.0.0 + - pip install helion==1.4.0 - pytest -v -s kernels/helion/ - label: Kernels Mamba Test # TBD diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 68e6a5762e..0259847d99 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -223,7 +223,7 @@ steps: - vllm/utils/import_utils.py - tests/kernels/helion/ commands: - - pip install helion==1.0.0 + - pip install helion==1.4.0 - pytest -v -s kernels/helion/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bd92ec9854..39de56e352 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ default_install_hook_types: default_stages: - pre-commit # Run locally - manual # Run in CI -exclude: '(vllm/third_party/.*|flash_qla/.*)' +exclude: '(vllm/third_party/.*|flash_qla/.*|csrc/sm70_turbomind/lmdeploy/.*)' repos: - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.14.0 @@ -35,6 +35,11 @@ repos: rev: v1.7.7 hooks: - id: actionlint +- repo: https://github.com/shellcheck-py/shellcheck-py + rev: 745eface02aef23e168a8afb6b5737818efbea95 # v0.11.0.1 + hooks: + - id: shellcheck + args: ["-s", "bash", "--format=gcc", "--severity=error"] - repo: https://github.com/astral-sh/uv-pre-commit rev: 0.11.1 hooks: @@ -177,11 +182,6 @@ repos: entry: python tools/pre_commit/mypy.py 1 "3.13" <<: *mypy_common stages: [manual] # Only run in CI - - id: shellcheck - name: Lint shell scripts - entry: tools/pre_commit/shellcheck.sh - language: script - types: [shell] - id: png-lint name: Lint PNG exports from excalidraw entry: tools/pre_commit/png-lint.sh diff --git a/CMakeLists.txt b/CMakeLists.txt index 21dc88d62a..c1fdaa5ea4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -397,6 +397,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") "${VLLM_SM70_TURBOMIND_LMDEPLOY_ROOT}/src/turbomind/kernels/gemm/kernel/sm70_884_16.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/tm_registry_sm70.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/fp8_qpn8_sm70.cu" + "${VLLM_SM70_TURBOMIND_ROOT}/ops/mxfp4_qpn_m1_sm70.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/nvfp4_qpn4_sm70.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/qwen38_prefill_cutlass.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/nvfp4_qpn2_sm70.cu" @@ -729,6 +730,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" AND SM70_TURBOMIND_ARCHS) SOURCES ${VLLM_SM70_SAMPLER_FRAGMENT_SRCS} COMPILE_FLAGS ${VLLM_GPU_FLAGS} ARCHITECTURES ${VLLM_GPU_ARCHES} + USE_SABI 3 WITH_SOABI) set_target_properties(_sm70_sampler_C PROPERTIES CXX_STANDARD 17 diff --git a/README.md b/README.md index f305d6e580..af2aba0f4d 100644 --- a/README.md +++ b/README.md @@ -349,9 +349,9 @@ python -m pip install -e . --no-build-isolation 请使用微信扫描下方二维码加入群组： -![1Cat-vLLM 微信交流群二维码](docs/assets/wechat-group-qr-7.png) +e52eca541f6c191928963c97ac1b38e5 -> 提示：微信群二维码通常 7 天内有效。若扫描失败或提示过期，请重新打开本页查看最新图片，或关注仓库更新。 +> 提示：微信群二维码通常 7 天内有效。若扫描失败或提示过期，请微信搜索微信号添加：YM_isi。 ## Repository Notes diff --git a/benchmarks/benchmark_dspark_api.py b/benchmarks/benchmark_dspark_api.py index 0279e4e896..a4cfa62140 100644 --- a/benchmarks/benchmark_dspark_api.py +++ b/benchmarks/benchmark_dspark_api.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Issue one reproducible OpenAI chat request for DSpark validation.""" diff --git a/benchmarks/benchmark_dspark_stream.py b/benchmarks/benchmark_dspark_stream.py index 16e7de9ae8..d40f1f237f 100644 --- a/benchmarks/benchmark_dspark_stream.py +++ b/benchmarks/benchmark_dspark_stream.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Measure one streamed completions request with separate TTFT and TPOT.""" diff --git a/benchmarks/benchmark_sm70_decode.py b/benchmarks/benchmark_sm70_decode.py index 19f27d6dea..93f4596100 100644 --- a/benchmarks/benchmark_sm70_decode.py +++ b/benchmarks/benchmark_sm70_decode.py @@ -52,20 +52,42 @@ def _sm70_fa2_d256_prefill_status(torch: Any) -> dict[str, Any]: "sm70_d256_splitd_n32_paged_fwd", ) optional_ops = ("sm70_d256_splitd_n32_dense_splitkv3_fwd",) - try: - importlib.import_module("vllm.vllm_flash_attn.flash_attn_interface") - except (AttributeError, ImportError, RuntimeError) as exc: - return { - "available": False, - "error": f"{type(exc).__name__}: {exc}", - "extension_file": None, - "extension_realpath": None, - "required_ops": {name: False for name in required_ops}, - "optional_ops": {name: False for name in optional_ops}, - } - extension_file = _module_file("vllm.vllm_flash_attn._vllm_fa2_C") namespace = getattr(torch.ops, "_vllm_fa2_C", None) + import_error: Exception | None = None + if namespace is None or not all(hasattr(namespace, name) for name in required_ops): + try: + importlib.import_module("vllm.vllm_flash_attn.flash_attn_interface") + except (AttributeError, ImportError, RuntimeError) as exc: + import_error = exc + namespace = getattr(torch.ops, "_vllm_fa2_C", None) + if namespace is None or not all(hasattr(namespace, name) for name in required_ops): + library_path = os.getenv("VLLM_SM70_FA2_D256_LIBRARY") + if library_path is not None: + try: + torch.ops.load_library(library_path) + except (OSError, RuntimeError) as load_exc: + return { + "available": False, + "error": f"{type(load_exc).__name__}: {load_exc}", + "extension_file": library_path, + "extension_realpath": str(Path(library_path).resolve()), + "required_ops": {name: False for name in required_ops}, + "optional_ops": {name: False for name in optional_ops}, + } + extension_file = library_path + elif import_error is not None: + return { + "available": False, + "error": f"{type(import_error).__name__}: {import_error}", + "extension_file": None, + "extension_realpath": None, + "required_ops": {name: False for name in required_ops}, + "optional_ops": {name: False for name in optional_ops}, + } + namespace = getattr(torch.ops, "_vllm_fa2_C", None) + if extension_file is None: + extension_file = os.getenv("VLLM_SM70_FA2_D256_LIBRARY") required_status = { name: namespace is not None and hasattr(namespace, name) for name in required_ops @@ -169,6 +191,7 @@ def _tracked_env() -> dict[str, str]: "VLLM_FLASH_", "FLASH_ATTN", "FLA_", + "PREFIX_", "CUDA_MODULE_LOADING", "TORCH_CUDA_ARCH_LIST", "TORCHINDUCTOR_", @@ -1172,6 +1195,8 @@ def _sm70_gdn_fla_policy() -> dict[str, Any]: def _sm70_moe_policy() -> dict[str, Any]: add_allreduce = _env_bool("VLLM_SM70_MOE_ADD_ALLREDUCE", False) + qwen38_router_topk = _env_bool("VLLM_SM70_QWEN38_ROUTER_TOPK", True) + qwen38_nvfp4_qpn_m1 = _env_bool("VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE", True) single_token_fastpath = _env_bool("VLLM_SM70_MOE_SINGLE_TOKEN_FASTPATH", False) single_token_permute = _env_bool( "VLLM_SM70_MOE_SINGLE_TOKEN_PERMUTE_FASTPATH", False @@ -1203,6 +1228,12 @@ def _sm70_moe_policy() -> dict[str, Any]: return { "VLLM_SM70_MOE_ADD_ALLREDUCE": os.environ.get("VLLM_SM70_MOE_ADD_ALLREDUCE"), "moe_add_allreduce_effective": add_allreduce, + "VLLM_SM70_QWEN38_ROUTER_TOPK": os.environ.get("VLLM_SM70_QWEN38_ROUTER_TOPK"), + "qwen38_router_topk_effective": qwen38_router_topk, + "VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE": os.environ.get( + "VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE" + ), + "qwen38_nvfp4_qpn_m1_effective": qwen38_nvfp4_qpn_m1, "VLLM_SM70_MOE_SINGLE_TOKEN_FASTPATH": os.environ.get( "VLLM_SM70_MOE_SINGLE_TOKEN_FASTPATH" ), @@ -1249,8 +1280,16 @@ def _sm70_moe_policy() -> dict[str, Any]: ), "profile_trace_effective": profile_trace, "expected_default": False, + "qwen38_router_topk_expected_default": True, + "qwen38_nvfp4_qpn_m1_expected_default": True, "single_token_unpermute_expected_default": True, "route_hit_oracle": ( + "The exact Qwen3.8 E512/K10 lane requires " + "qwen38_router_topk_effective=true and the " + "`SM70 Qwen3.8 E512/K10 router top-k path enabled` marker. " + "The accepted routed-expert lane additionally requires " + "qwen38_nvfp4_qpn_m1_effective=true and the `SM70 Qwen3.8 " + "NVFP4 direct QPN-M1 expert path enabled` marker. " "For CUDA-graph fast-path acceptance, require the C++ trace " "`SM70 custom all_reduce_sum2 op reached ... capture=active`; " "the Python MoERunner candidate log alone is not a custom-op hit. " @@ -1267,6 +1306,7 @@ def _sm70_moe_policy() -> dict[str, Any]: def _sm70_sampling_policy() -> dict[str, Any]: + use_v2_model_runner = _env_bool("VLLM_USE_V2_MODEL_RUNNER", False) flash_0dot3_compile = _env_bool( "VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH", False, @@ -1279,6 +1319,8 @@ def _sm70_sampling_policy() -> dict[str, Any]: full_lm_head = _env_bool("VLLM_SM70_ENABLE_LM_HEAD_FASTPATH", False) dense_f16 = _env_bool("VLLM_SM70_ENABLE_DENSE_F16_FASTPATH", False) return { + "VLLM_USE_V2_MODEL_RUNNER": os.environ.get("VLLM_USE_V2_MODEL_RUNNER"), + "use_v2_model_runner_effective": use_v2_model_runner, "VLLM_SM70_GREEDY_TOKEN_FASTPATH": os.environ.get( "VLLM_SM70_GREEDY_TOKEN_FASTPATH" ), @@ -1319,6 +1361,14 @@ def _sm70_sampling_policy() -> dict[str, Any]: and not full_lm_head and not custom_top1_ar ), + "v2_greedy_tp_local_pair_policy": ( + use_v2_model_runner + and greedy_fastpath + and not lm_head_top1 + and not lm_head_top1_tc + and not full_lm_head + and not custom_top1_ar + ), "route_hit_oracle": ( "Accepted default greedy decode requires " "`SM70 LM head top1 layout prepared` plus " @@ -1326,7 +1376,10 @@ def _sm70_sampling_policy() -> dict[str, Any]: "The SM70 Flash-V100 compile-graph quality lane instead requires " "VLLM_SM70_LM_HEAD_TOP1=0 and " "compile_graph_local_logits_top1_policy=true, preserving greedy " - "pair-gather while avoiding the fused top1 epilogue. " + "pair-gather while avoiding the fused top1 epilogue. The MRv2 " + "single-token lane additionally requires " + "v2_greedy_tp_local_pair_policy=true and the " + "`SM70 MRv2 greedy TP-local pair path enabled` log marker. " "`SM70 dense fp16 fast path enabled for LM head` is expected only " "when the explicit full LM-head fast path gate is enabled. " "The 0.0.3 LM-head baseline lane requires full LM-head fast path " diff --git a/benchmarks/benchmark_sm70_dflash2_gsm8k.py b/benchmarks/benchmark_sm70_dflash2_gsm8k.py index ac92519225..47f96a1847 100644 --- a/benchmarks/benchmark_sm70_dflash2_gsm8k.py +++ b/benchmarks/benchmark_sm70_dflash2_gsm8k.py @@ -66,6 +66,12 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--max-num-batched-tokens", type=int, default=512) parser.add_argument("--max-num-seqs", type=int, default=4) parser.add_argument("--sequential", action="store_true") + parser.add_argument("--enable-prefix-caching", action="store_true") + parser.add_argument( + "--mamba-cache-mode", + choices=("all", "align", "none"), + help="Override the engine Mamba cache mode for practical-contract runs.", + ) parser.add_argument("--gpu-memory-utilization", type=float, default=0.8) parser.add_argument( "--target-kv-cache-dtype", @@ -85,12 +91,26 @@ def _parse_args() -> argparse.Namespace: help="Draft-only attention backend; the target remains on FLASH_ATTN_V100.", ) parser.add_argument("--seed", type=int, default=0) - parser.add_argument( + seed_group = parser.add_mutually_exclusive_group() + seed_group.add_argument( "--request-seed", type=int, default=0, help="Sampling seed for every request; use -1 for server-style random seeds.", ) + seed_group.add_argument( + "--request-seeds", + help=( + "Comma-separated non-negative sampling seeds. Each seed replays the " + "selected prompt set without reloading the model." + ), + ) + parser.add_argument( + "--request-seed-mode", + choices=("fixed", "index"), + default="fixed", + help="Use each seed as-is or add the original dataset index per request.", + ) parser.add_argument( "--cuda-profiler-capture", action="store_true", @@ -107,6 +127,23 @@ def _sha256_file(path: Path) -> str: return digest.hexdigest() +def _request_seeds(args: argparse.Namespace) -> list[int | None]: + if args.request_seeds is None: + if args.request_seed < -1: + raise ValueError("--request-seed must be -1 or a non-negative integer") + return [None if args.request_seed == -1 else args.request_seed] + + try: + seeds = [int(value.strip()) for value in args.request_seeds.split(",")] + except ValueError as exc: + raise ValueError("--request-seeds must be comma-separated integers") from exc + if not seeds or any(seed < 0 for seed in seeds): + raise ValueError("--request-seeds must contain non-negative integers") + if len(set(seeds)) != len(seeds): + raise ValueError("--request-seeds must not contain duplicates") + return seeds + + def _model_weight_files(model: Path) -> dict[str, str]: index = model / "model.safetensors.index.json" if not index.is_file(): @@ -219,10 +256,11 @@ def main() -> int: args = _parse_args() if args.num_questions <= 0: raise ValueError("--num-questions must be positive") - if args.request_seed < -1: - raise ValueError("--request-seed must be -1 or a non-negative integer") if args.mode == "dflash" and args.draft_model is None: raise ValueError("--draft-model is required for --mode dflash") + if args.request_seed_mode == "index" and not args.sequential: + raise ValueError("--request-seed-mode index requires --sequential") + request_seeds = _request_seeds(args) import torch import vllm._C as vllm_c @@ -275,12 +313,14 @@ def main() -> int: "max_num_batched_tokens": args.max_num_batched_tokens, "max_num_seqs": args.max_num_seqs, "gpu_memory_utilization": args.gpu_memory_utilization, - "enable_prefix_caching": False, + "enable_prefix_caching": args.enable_prefix_caching, "disable_log_stats": False, "enforce_eager": args.enforce_eager, "seed": args.seed, "speculative_config": speculative_config, } + if args.mamba_cache_mode is not None: + engine_kwargs["mamba_cache_mode"] = args.mamba_cache_mode if args.cuda_profiler_capture: engine_kwargs["profiler_config"] = {"profiler": "cuda"} @@ -288,42 +328,84 @@ def main() -> int: llm = LLM(**engine_kwargs) load_seconds = time.perf_counter() - load_started - request_seed = None if args.request_seed == -1 else args.request_seed + warmup_seed = request_seeds[0] + if warmup_seed is not None and args.request_seed_mode == "index": + warmup_seed += rows[0][0] warmup_sampling = SamplingParams( temperature=args.temperature, top_p=0.95, top_k=20, max_tokens=args.warmup_tokens, - seed=request_seed, + seed=warmup_seed, skip_special_tokens=False, ) llm.generate([prompts[0]], warmup_sampling, use_tqdm=False) - sampling = SamplingParams( - temperature=args.temperature, - top_p=0.95, - top_k=20, - max_tokens=args.max_tokens, - seed=request_seed, - skip_special_tokens=False, - ) spec_before = _spec_metrics_snapshot(llm) if args.cuda_profiler_capture: llm.start_profile() started = time.perf_counter() - if args.sequential: - outputs = [] - request_spec_metrics = [] - for prompt in prompts: - request_spec_before = _spec_metrics_snapshot(llm) - outputs.append(llm.generate([prompt], sampling, use_tqdm=False)[0]) - request_spec_after = _spec_metrics_snapshot(llm) - request_spec_metrics.append( - _diff_spec_metrics(request_spec_before, request_spec_after) + outputs = [] + request_spec_metrics = [] + case_inputs = [] + for seed_base in request_seeds: + if args.sequential: + seed_outputs = [] + seed_spec_metrics = [] + actual_seeds = [] + for (dataset_index, _row), prompt in zip(rows, prompts, strict=True): + request_seed = seed_base + if request_seed is not None and args.request_seed_mode == "index": + request_seed += dataset_index + sampling = SamplingParams( + temperature=args.temperature, + top_p=0.95, + top_k=20, + max_tokens=args.max_tokens, + seed=request_seed, + skip_special_tokens=False, + ) + request_spec_before = _spec_metrics_snapshot(llm) + seed_outputs.append(llm.generate([prompt], sampling, use_tqdm=False)[0]) + request_spec_after = _spec_metrics_snapshot(llm) + seed_spec_metrics.append( + _diff_spec_metrics(request_spec_before, request_spec_after) + ) + actual_seeds.append(request_seed) + else: + sampling = SamplingParams( + temperature=args.temperature, + top_p=0.95, + top_k=20, + max_tokens=args.max_tokens, + seed=seed_base, + skip_special_tokens=False, + ) + seed_outputs = llm.generate(prompts, sampling, use_tqdm=False) + seed_spec_metrics = [None] * len(seed_outputs) + actual_seeds = [seed_base] * len(seed_outputs) + outputs.extend(seed_outputs) + request_spec_metrics.extend(seed_spec_metrics) + case_inputs.extend( + ( + seed_base, + request_seed, + dataset_index, + row, + prompt_content, + prompt_tokens, + ) + for request_seed, ( + dataset_index, + row, + ), prompt_content, prompt_tokens in zip( + actual_seeds, + rows, + prompt_contents, + prompt_token_counts, + strict=True, ) - else: - outputs = llm.generate(prompts, sampling, use_tqdm=False) - request_spec_metrics = [None] * len(outputs) + ) elapsed_seconds = time.perf_counter() - started if args.cuda_profiler_capture: llm.stop_profile() @@ -331,10 +413,17 @@ def main() -> int: cases = [] for ( + request_seed_base, + request_seed, dataset_index, row, - ), prompt_content, prompt_tokens, output, request_spec in zip( - rows, prompt_contents, prompt_token_counts, outputs, request_spec_metrics + prompt_content, + prompt_tokens, + ), output, request_spec in zip( + case_inputs, + outputs, + request_spec_metrics, + strict=True, ): result = output.outputs[0] token_ids = list(result.token_ids) @@ -348,7 +437,11 @@ def main() -> int: correct = None cases.append( { + "request_seed_base": request_seed_base, + "request_seed": request_seed, "dataset_index": dataset_index, + "suite": row.get("_suite"), + "suite_index": row.get("_suite_index"), "question": row.get("question"), "prompt_content": prompt_content, "expected_answer": expected, @@ -400,6 +493,7 @@ def main() -> int: "dataset_format": args.dataset_format, "start_index": args.start_index, "num_questions": args.num_questions, + "num_samples": len(cases), "dataset_order": args.dataset_order, "model": str(args.model), "model_config_sha256": _sha256_file(args.model / "config.json"), @@ -418,7 +512,14 @@ def main() -> int: "top_p": 0.95, "top_k": 20, "max_tokens": args.max_tokens, - "seed": request_seed, + "seed": ( + request_seeds[0] + if len(request_seeds) == 1 and args.request_seed_mode == "fixed" + else None + ), + "seeds": (request_seeds if args.request_seed_mode == "fixed" else None), + "seed_bases": request_seeds, + "seed_mode": args.request_seed_mode, "ignore_eos": False, "thinking": True, "reasoning_effort": "xhigh", diff --git a/benchmarks/benchmark_sm70_flashinfer_fixed_entry.py b/benchmarks/benchmark_sm70_flashinfer_fixed_entry.py index 2482003308..4eb48de306 100644 --- a/benchmarks/benchmark_sm70_flashinfer_fixed_entry.py +++ b/benchmarks/benchmark_sm70_flashinfer_fixed_entry.py @@ -263,9 +263,7 @@ def _active_ctas_per_sm( threads_per_cta: int, shared_bytes_per_cta: int, ) -> int: - register_limited = SM70_REGISTERS_PER_SM // ( - registers_per_thread * threads_per_cta - ) + register_limited = SM70_REGISTERS_PER_SM // (registers_per_thread * threads_per_cta) shared_limited = SM70_SHARED_BYTES_PER_SM // shared_bytes_per_cta thread_limited = SM70_THREADS_PER_SM // threads_per_cta return min(register_limited, shared_limited, thread_limited, SM70_MAX_CTA_PER_SM) @@ -350,9 +348,7 @@ def _paged_resource_gate(kernel: Mapping[str, Any]) -> dict[str, Any]: f"requires <= {PAGED_MAX_SHARED_BYTES}" ) if kernel["active_ctas_per_sm"] != 2: - reasons.append( - f"active_ctas_per_sm={kernel['active_ctas_per_sm']}, requires 2" - ) + reasons.append(f"active_ctas_per_sm={kernel['active_ctas_per_sm']}, requires 2") instructions = kernel["instructions"] if instructions["hmma"] == 0: reasons.append("paged target has no HMMA") @@ -617,15 +613,9 @@ def _pairwise_stats( return { "right_minus_left_mean_ms": float(statistics.mean(deltas)), "right_minus_left_p50_ms": float(statistics.median(deltas)), - "right_faster_wins": sum( - right_ms < left_ms for left_ms, right_ms in paired - ), - "left_faster_wins": sum( - left_ms < right_ms for left_ms, right_ms in paired - ), - "ties": sum( - left_ms == right_ms for left_ms, right_ms in paired - ), + "right_faster_wins": sum(right_ms < left_ms for left_ms, right_ms in paired), + "left_faster_wins": sum(left_ms < right_ms for left_ms, right_ms in paired), + "ties": sum(left_ms == right_ms for left_ms, right_ms in paired), "right_vs_left_mean_pct": ( (right_stats["mean_ms"] / left_stats["mean_ms"] - 1.0) * 100.0 ), diff --git a/benchmarks/benchmark_sm70_flashqla_indexed_state_cache.py b/benchmarks/benchmark_sm70_flashqla_indexed_state_cache.py index f83336bfbc..e94a7f9ace 100644 --- a/benchmarks/benchmark_sm70_flashqla_indexed_state_cache.py +++ b/benchmarks/benchmark_sm70_flashqla_indexed_state_cache.py @@ -72,27 +72,38 @@ def main() -> None: dtype = torch.float16 scale = args.dim**-0.5 - q = (torch.randn(1, args.tokens, args.q_heads, args.dim, - device=device, dtype=dtype) * 0.05).contiguous() + q = ( + torch.randn(1, args.tokens, args.q_heads, args.dim, device=device, dtype=dtype) + * 0.05 + ).contiguous() k = (torch.randn_like(q) * 0.05).contiguous() - v = (torch.randn(1, args.tokens, args.v_heads, args.dim, - device=device, dtype=dtype) * 0.1).contiguous() - g = (torch.randn(1, args.tokens, args.v_heads, - device=device, dtype=torch.float32) * 0.02 - 0.04).contiguous() - beta = torch.rand(1, args.tokens, args.v_heads, - device=device, dtype=torch.float32).contiguous() + v = ( + torch.randn(1, args.tokens, args.v_heads, args.dim, device=device, dtype=dtype) + * 0.1 + ).contiguous() + g = ( + torch.randn(1, args.tokens, args.v_heads, device=device, dtype=torch.float32) + * 0.02 + - 0.04 + ).contiguous() + beta = torch.rand( + 1, args.tokens, args.v_heads, device=device, dtype=torch.float32 + ).contiguous() cu_seqlens = torch.tensor([0, args.tokens], device=device, dtype=torch.int32) state_indices = torch.tensor([args.state_index], device=device, dtype=torch.long) has_initial_state = torch.ones(1, device=device, dtype=torch.bool) - base_state_cache = (torch.randn( - args.state_slots, - args.v_heads, - args.dim, - args.dim, - device=device, - dtype=torch.float32, - ) * 0.01).contiguous() + base_state_cache = ( + torch.randn( + args.state_slots, + args.v_heads, + args.dim, + args.dim, + device=device, + dtype=torch.float32, + ) + * 0.01 + ).contiguous() gather_cache = base_state_cache.clone() indexed_cache = base_state_cache.clone() gather_out = torch.empty_like(v) @@ -151,8 +162,9 @@ def indexed_call(): correctness = { "out": _diff(gather_result, indexed_result), - "state_slot": _diff(gather_cache[args.state_index], - indexed_cache[args.state_index]), + "state_slot": _diff( + gather_cache[args.state_index], indexed_cache[args.state_index] + ), } if args.state_slots > 1: untouched_index = 0 if args.state_index != 0 else args.state_slots - 1 diff --git a/benchmarks/benchmark_sm70_fp8_fixed_vs_dynamic.py b/benchmarks/benchmark_sm70_fp8_fixed_vs_dynamic.py index 7d21e6def3..fd9e069e54 100644 --- a/benchmarks/benchmark_sm70_fp8_fixed_vs_dynamic.py +++ b/benchmarks/benchmark_sm70_fp8_fixed_vs_dynamic.py @@ -65,8 +65,9 @@ def _candidate_layers(model_path: Path) -> tuple[dict[str, str], list[str]]: return weight_map, layers -def _select_layers(layers: list[str], layer_ids: set[int], - max_layers: int) -> list[str]: +def _select_layers( + layers: list[str], layer_ids: set[int], max_layers: int +) -> list[str]: selected: list[str] = [] seen_kind: set[str] = set() for prefix in layers: @@ -100,8 +101,7 @@ def _run_case( "skip": "shape_not_multiple_128", } - tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare( - qweight, scales, group_size) + tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare(qweight, scales, group_size) x = torch.randn((m, k), device=qweight.device, dtype=torch.float16) fixed = torch.empty((m, n), device=qweight.device, dtype=torch.float16) @@ -120,8 +120,7 @@ def _run_case( os.environ["VLLM_SM70_FP8_TUNE_SMALL_SHAPES"] = "0" fixed_after = torch.empty((m, n), device=qweight.device, dtype=torch.float16) - sm70_ops.fp8_gemm_sm70_out_meta( - fixed_after, x, tm_weight, tm_scales, meta) + sm70_ops.fp8_gemm_sm70_out_meta(fixed_after, x, tm_weight, tm_scales, meta) torch.accelerator.synchronize(qweight.device) diff_dynamic = (fixed.float() - dynamic.float()).abs() @@ -178,8 +177,7 @@ def main() -> int: results: list[dict[str, Any]] = [] try: for prefix in selected: - qweight, scales = _load_fp8_layer( - args.model, weight_map, prefix, device) + qweight, scales = _load_fp8_layer(args.model, weight_map, prefix, device) for m in args.m: results.append(_run_case(prefix, qweight, scales, m)) finally: @@ -189,7 +187,8 @@ def main() -> int: os.environ["VLLM_SM70_FP8_TUNE_SMALL_SHAPES"] = original_env failures = [ - item for item in results + item + for item in results if item.get("fixed_vs_dynamic_max_diff", 0.0) != 0.0 or not item.get("fixed_vs_dynamic_equal", True) ] diff --git a/benchmarks/benchmark_sm70_hmma_c2a_mincomm.py b/benchmarks/benchmark_sm70_hmma_c2a_mincomm.py index 8eed68503a..1e9baafc70 100644 --- a/benchmarks/benchmark_sm70_hmma_c2a_mincomm.py +++ b/benchmarks/benchmark_sm70_hmma_c2a_mincomm.py @@ -7,13 +7,14 @@ import argparse import json import os -import re import shutil import subprocess import tempfile from pathlib import Path from typing import Any +import regex as re + SYMBOLS = { "reference": "sm70_hmma_c2a_reference", "candidate": "sm70_hmma_c2a_mincomm", diff --git a/benchmarks/benchmark_sm70_instruction_latency.py b/benchmarks/benchmark_sm70_instruction_latency.py index 05ec0612f7..d75ddb1397 100644 --- a/benchmarks/benchmark_sm70_instruction_latency.py +++ b/benchmarks/benchmark_sm70_instruction_latency.py @@ -7,13 +7,14 @@ import argparse import json import os -import re import shutil import subprocess import tempfile from pathlib import Path from typing import Any +import regex as re + def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() diff --git a/benchmarks/benchmark_sm70_lm_head_top1.py b/benchmarks/benchmark_sm70_lm_head_top1.py index a34e3e0fc1..4e846e1c50 100644 --- a/benchmarks/benchmark_sm70_lm_head_top1.py +++ b/benchmarks/benchmark_sm70_lm_head_top1.py @@ -111,9 +111,7 @@ def main() -> int: x = torch.randn((args.m, args.k), device=device, dtype=torch.float16) weight = torch.randn((args.n, args.k), device=device, dtype=torch.float16) - torch_logits_out = torch.empty( - (args.m, args.n), device=device, dtype=torch.float16 - ) + torch_logits_out = torch.empty((args.m, args.n), device=device, dtype=torch.float16) prepared = sm70_ops.sm70_f16_prepare(weight) tm_weight = prepared[0] k_ld = int(prepared[1][0].item()) @@ -121,7 +119,7 @@ def main() -> int: def torch_mm_max() -> tuple[torch.Tensor, torch.Tensor]: logits = torch.mm(x, weight.t()) if args.pad: - logits[:, -args.pad:] = -float("inf") + logits[:, -args.pad :] = -float("inf") values, indices = logits.max(dim=-1) return values, indices.to(torch.int64) + args.vocab_start @@ -189,15 +187,21 @@ def tensor_core_top1() -> None: def sm70_gemm_max() -> tuple[torch.Tensor, torch.Tensor]: sm70_ops.sm70_f16_gemm_out(logits_out, x, tm_weight, k_ld, False) if args.pad: - logits_out[:, -args.pad:] = -float("inf") + logits_out[:, -args.pad :] = -float("inf") values, indices = logits_out.max(dim=-1) return values, indices.to(torch.int64) + args.vocab_start sm70_values, sm70_indices = sm70_gemm_max() sm70_ms = _bench_cuda_ms(lambda: sm70_gemm_max(), args.warmup, args.iters) results.append( - _stats("sm70_gemm_then_max", sm70_values, sm70_indices, ref_values, - ref_indices, sm70_ms) + _stats( + "sm70_gemm_then_max", + sm70_values, + sm70_indices, + ref_values, + ref_indices, + sm70_ms, + ) ) torch_mm_ms = _bench_cuda_ms(lambda: torch_mm_max(), args.warmup, args.iters) diff --git a/benchmarks/benchmark_sm70_long_quality_suffix.py b/benchmarks/benchmark_sm70_long_quality_suffix.py index 0a9c06a273..f57a802ce3 100644 --- a/benchmarks/benchmark_sm70_long_quality_suffix.py +++ b/benchmarks/benchmark_sm70_long_quality_suffix.py @@ -220,9 +220,9 @@ def main() -> None: "im_end_count": result["im_end_count"], "unique_tokens": result["unique_tokens"], "top_token_count_id": result["top_token_count_id"], - "steady_decode_tps": ( - result["request_metrics"] or {} - ).get("steady_decode_tps"), + "steady_decode_tps": (result["request_metrics"] or {}).get( + "steady_decode_tps" + ), "prefix": text[:160], }, ensure_ascii=False, diff --git a/benchmarks/benchmark_sm70_model_tokens.py b/benchmarks/benchmark_sm70_model_tokens.py index 2f0d580bcf..75112d51fd 100644 --- a/benchmarks/benchmark_sm70_model_tokens.py +++ b/benchmarks/benchmark_sm70_model_tokens.py @@ -240,6 +240,66 @@ def _spec_decoding_summary( } +def _spec_decoding_delta( + before: list[dict[str, Any]], + after: list[dict[str, Any]], +) -> dict[str, Any] | None: + """Return speculative-decoding counters added between two snapshots.""" + + before_summary = _spec_decoding_summary(before) + after_summary = _spec_decoding_summary(after) + if after_summary is None: + return None + + before_drafts = int(before_summary["num_drafts"]) if before_summary else 0 + before_draft_tokens = ( + int(before_summary["num_draft_tokens"]) if before_summary else 0 + ) + before_accepted = ( + int(before_summary["num_accepted_tokens"]) if before_summary else 0 + ) + after_drafts = int(after_summary["num_drafts"]) + after_draft_tokens = int(after_summary["num_draft_tokens"]) + after_accepted = int(after_summary["num_accepted_tokens"]) + + num_drafts = after_drafts - before_drafts + num_draft_tokens = after_draft_tokens - before_draft_tokens + num_accepted_tokens = after_accepted - before_accepted + if num_drafts <= 0: + return None + if num_draft_tokens < 0 or num_accepted_tokens < 0: + raise RuntimeError("Speculative-decoding counters moved backwards") + + before_per_pos = ( + list(before_summary["accepted_tokens_per_pos"]) if before_summary else [] + ) + after_per_pos = list(after_summary["accepted_tokens_per_pos"]) + width = max(len(before_per_pos), len(after_per_pos)) + before_per_pos.extend([0] * (width - len(before_per_pos))) + after_per_pos.extend([0] * (width - len(after_per_pos))) + accepted_tokens_per_pos = [ + end - start for start, end in zip(before_per_pos, after_per_pos, strict=True) + ] + if any(value < 0 for value in accepted_tokens_per_pos): + raise RuntimeError("Per-position speculative counters moved backwards") + + avg_accepted_tokens = num_accepted_tokens / num_drafts + return { + "num_drafts": num_drafts, + "num_draft_tokens": num_draft_tokens, + "num_accepted_tokens": num_accepted_tokens, + "avg_accepted_tokens_no_bonus": avg_accepted_tokens, + "mean_acceptance_length": 1 + avg_accepted_tokens, + "draft_acceptance_rate": ( + num_accepted_tokens / num_draft_tokens if num_draft_tokens else None + ), + "accepted_tokens_per_pos": accepted_tokens_per_pos, + "per_position_acceptance_rate": [ + value / num_drafts for value in accepted_tokens_per_pos + ], + } + + def _load_prompts(args: argparse.Namespace) -> list[Any]: if args.input_lens is not None: if args.input_len is not None or args.prompt or args.prompts_json is not None: @@ -1713,6 +1773,7 @@ def _dump(args: argparse.Namespace) -> int: torch.accelerator.synchronize() torch.cuda.cudart().cudaProfilerStart() generate_seconds_by_repeat: list[float] = [] + sequential_prompt_metrics: list[dict[str, Any]] = [] outputs = [] try: for repeat_index in range(args.repeat_count): @@ -1724,8 +1785,24 @@ def _dump(args: argparse.Namespace) -> int: raise RuntimeError("Failed to reset the idle prefix cache") generate_start = time.perf_counter() if args.sequential_prompts: - for prompt in prompts: + metrics_before_prompt = _metric_snapshot(llm) + for prompt_index, prompt in enumerate(prompts): + prompt_start = time.perf_counter() outputs.extend(llm.generate([prompt], sampling_params)) + prompt_seconds = time.perf_counter() - prompt_start + metrics_after_prompt = _metric_snapshot(llm) + sequential_prompt_metrics.append( + { + "repeat_index": repeat_index, + "prompt_index": prompt_index, + "generate_seconds": prompt_seconds, + "spec_decoding_metrics": _spec_decoding_delta( + metrics_before_prompt, + metrics_after_prompt, + ), + } + ) + metrics_before_prompt = metrics_after_prompt else: outputs.extend(llm.generate(prompts, sampling_params)) generate_seconds_by_repeat.append(time.perf_counter() - generate_start) @@ -1831,6 +1908,7 @@ def _dump(args: argparse.Namespace) -> int: "eos_token_ids": eos_token_ids, "metrics_snapshot": metrics_snapshot, "spec_decoding_metrics": _spec_decoding_summary(metrics_snapshot), + "sequential_prompt_metrics": sequential_prompt_metrics, "sampling_params": { "max_tokens": args.max_tokens, "temperature": args.temperature, diff --git a/benchmarks/benchmark_sm70_native_bm32_allp_crossblock_micro.py b/benchmarks/benchmark_sm70_native_bm32_allp_crossblock_micro.py index 10771b3a8b..6069cf2926 100644 --- a/benchmarks/benchmark_sm70_native_bm32_allp_crossblock_micro.py +++ b/benchmarks/benchmark_sm70_native_bm32_allp_crossblock_micro.py @@ -103,8 +103,7 @@ def _require_environment(args: argparse.Namespace) -> None: raise ValueError(f"formal runs require exactly {K_DEFAULT_ROUNDS} rounds.") if args.launches_per_sample != K_DEFAULT_LAUNCHES: raise ValueError( - "formal runs require exactly " - f"{K_DEFAULT_LAUNCHES} launches per sample." + f"formal runs require exactly {K_DEFAULT_LAUNCHES} launches per sample." ) @@ -261,12 +260,8 @@ def _inspect_sass(binary: Path) -> dict[str, Any]: "kernels": kernels, "static_instruction_delta": { mnemonic: ( - kernels.get("candidate", {}) - .get("instructions", {}) - .get(mnemonic, 0) - - kernels.get("baseline", {}) - .get("instructions", {}) - .get(mnemonic, 0) + kernels.get("candidate", {}).get("instructions", {}).get(mnemonic, 0) + - kernels.get("baseline", {}).get("instructions", {}).get(mnemonic, 0) ) for mnemonic in ("hmma", "ldg", "lds", "sts", "ldl", "stl", "bar") }, @@ -302,13 +297,11 @@ def _resource_gate( ) if resource.get("registers_per_thread", 65) > 64: reasons.append( - f"{path} REG={resource.get('registers_per_thread')}, " - "requires <=64" + f"{path} REG={resource.get('registers_per_thread')}, requires <=64" ) if resource.get("local_bytes_per_thread") != 0: reasons.append( - f"{path} LOCAL={resource.get('local_bytes_per_thread')}, " - "requires 0" + f"{path} LOCAL={resource.get('local_bytes_per_thread')}, requires 0" ) if resource.get("static_shared_bytes") != K_ALLP_SHARED_BYTES: reasons.append( @@ -333,8 +326,7 @@ def _resource_gate( continue if ptxas["registers_per_thread"] > 64: reasons.append( - f"PTXAS {path} REG={ptxas['registers_per_thread']}, " - "requires <=64" + f"PTXAS {path} REG={ptxas['registers_per_thread']}, requires <=64" ) for key in ("stack_frame_bytes", "spill_store_bytes", "spill_load_bytes"): if ptxas[key] != 0: @@ -356,9 +348,7 @@ def _resource_gate( instructions = path_sass.get("instructions", {}) for mnemonic in ("hmma", "ldg", "lds", "sts"): if instructions.get(mnemonic, 0) == 0: - reasons.append( - f"{path} SASS has no {mnemonic.upper()} instruction" - ) + reasons.append(f"{path} SASS has no {mnemonic.upper()} instruction") if instructions.get("ldl", 0) or instructions.get("stl", 0): reasons.append( f"{path} SASS has local instructions " @@ -397,40 +387,31 @@ def _resource_gate( return not reasons, reasons -def _correctness_gate( - payload: dict[str, Any], groups: int -) -> tuple[bool, list[str]]: +def _correctness_gate(payload: dict[str, Any], groups: int) -> tuple[bool, list[str]]: reasons: list[str] = [] exactness = payload.get("exactness") or {} expected_words = groups * K_M * K_D // 2 if exactness.get("word_dtype") != "uint32 packed fp16": reasons.append( - "word_dtype=" - f"{exactness.get('word_dtype')}, requires uint32 packed fp16" + f"word_dtype={exactness.get('word_dtype')}, requires uint32 packed fp16" ) if exactness.get("word_count") != expected_words: reasons.append( - "word_count=" - f"{exactness.get('word_count')}, requires {expected_words}" + f"word_count={exactness.get('word_count')}, requires {expected_words}" ) if exactness.get("full_output") is not True: reasons.append("C++ probe did not compare the full output") if exactness.get("bitwise_equal") is not True: reasons.append("candidate output is not bitwise equal to baseline") if exactness.get("mismatch_words") != 0: - reasons.append( - "mismatch_words=" - f"{exactness.get('mismatch_words')}, requires 0" - ) + reasons.append(f"mismatch_words={exactness.get('mismatch_words')}, requires 0") xor = exactness.get("xor", {}) if xor.get("max_word") != 0 or xor.get("reduction") != 0: reasons.append("packed uint32 XOR result is nonzero") return not reasons, reasons -def _performance_gate( - payload: dict[str, Any], rounds: int -) -> tuple[bool, list[str]]: +def _performance_gate(payload: dict[str, Any], rounds: int) -> tuple[bool, list[str]]: reasons: list[str] = [] pairs = payload.get("pairs") or {} if pairs.get("count") != rounds: @@ -441,13 +422,10 @@ def _performance_gate( f"candidate_faster={pairs.get('candidate_faster')}, " f"requires >= {required_wins}/{rounds}" ) - speedup_pct = (payload.get("timing") or {}).get( - "candidate_speedup_vs_baseline_pct" - ) + speedup_pct = (payload.get("timing") or {}).get("candidate_speedup_vs_baseline_pct") if speedup_pct is None or speedup_pct < K_REQUIRED_SPEEDUP_PCT: reasons.append( - f"median speedup={speedup_pct}, requires " - f">= {K_REQUIRED_SPEEDUP_PCT}%" + f"median speedup={speedup_pct}, requires >= {K_REQUIRED_SPEEDUP_PCT}%" ) return not reasons, reasons @@ -666,14 +644,11 @@ def _run_ncu_kernel( "mio_throttle_per_warp_active": metrics.get( "smsp__warp_issue_stalled_mio_throttle_per_warp_active" ), - "tensor_instructions": metrics.get( - "smsp__inst_executed_pipe_tensor.sum" - ), + "tensor_instructions": metrics.get("smsp__inst_executed_pipe_tensor.sum"), }, "stdout_tail": completed.stdout[-3000:], "stderr_tail": completed.stderr[-3000:], - "permission_error": "ERR_NVGPUCTRPERM" - in (completed.stdout + completed.stderr), + "permission_error": "ERR_NVGPUCTRPERM" in (completed.stdout + completed.stderr), } @@ -830,8 +805,7 @@ def _main() -> int: "benchmark": "sm70_native_bm32_allp_crossblock_micro", "target": "sm_70", "matrix": [ - {"pattern": pattern, "nblocks": nblocks} - for pattern, nblocks in K_CASES + {"pattern": pattern, "nblocks": nblocks} for pattern, nblocks in K_CASES ], "configuration": { "groups": args.groups, diff --git a/benchmarks/benchmark_sm70_native_bm32_allp_scratch_micro.py b/benchmarks/benchmark_sm70_native_bm32_allp_scratch_micro.py index 4679eff592..177a128420 100644 --- a/benchmarks/benchmark_sm70_native_bm32_allp_scratch_micro.py +++ b/benchmarks/benchmark_sm70_native_bm32_allp_scratch_micro.py @@ -94,10 +94,7 @@ def _require_environment(args: argparse.Namespace) -> None: if args.rounds < K_DEFAULT_ROUNDS: raise ValueError(f"rounds must be at least {K_DEFAULT_ROUNDS}.") if args.launches_per_sample < K_DEFAULT_LAUNCHES: - raise ValueError( - "launches-per-sample must be at least " - f"{K_DEFAULT_LAUNCHES}." - ) + raise ValueError(f"launches-per-sample must be at least {K_DEFAULT_LAUNCHES}.") def _graphics_clock_mhz(physical_gpu: int) -> int: @@ -248,8 +245,7 @@ def _inspect_sass(binary: Path) -> dict[str, Any]: "ctas_per_group": ctas_per_group, "instructions": instructions, "instructions_per_group": { - name: count * ctas_per_group - for name, count in instructions.items() + name: count * ctas_per_group for name, count in instructions.items() }, } baseline_instructions = kernels.get("baseline", {}).get("instructions", {}) @@ -293,13 +289,11 @@ def _resource_gate( ) if resource.get("registers_per_thread", 65) > 64: reasons.append( - f"{path} REG={resource.get('registers_per_thread')}, " - "requires <=64" + f"{path} REG={resource.get('registers_per_thread')}, requires <=64" ) if resource.get("local_bytes_per_thread") != 0: reasons.append( - f"{path} LOCAL={resource.get('local_bytes_per_thread')}, " - "requires 0" + f"{path} LOCAL={resource.get('local_bytes_per_thread')}, requires 0" ) if resource.get("static_shared_bytes") != K_ALLP_SHARED_BYTES: reasons.append( @@ -326,8 +320,7 @@ def _resource_gate( continue if ptxas["registers_per_thread"] > 64: reasons.append( - f"PTXAS {path} REG={ptxas['registers_per_thread']}, " - "requires <=64" + f"PTXAS {path} REG={ptxas['registers_per_thread']}, requires <=64" ) for key in ("stack_frame_bytes", "spill_store_bytes", "spill_load_bytes"): if ptxas[key] != 0: @@ -348,9 +341,7 @@ def _resource_gate( instructions = path_sass.get("instructions", {}) for mnemonic in ("hmma", "ldg", "lds"): if instructions.get(mnemonic, 0) == 0: - reasons.append( - f"{path} SASS has no {mnemonic.upper()} instruction" - ) + reasons.append(f"{path} SASS has no {mnemonic.upper()} instruction") if instructions.get("ldl", 0) or instructions.get("stl", 0): reasons.append( f"{path} SASS has local instructions " @@ -380,46 +371,35 @@ def _resource_gate( } for key, expected in expected_scratch.items(): if scratch.get(key) != expected: - reasons.append( - f"scratch {key}={scratch.get(key)}, requires {expected}" - ) + reasons.append(f"scratch {key}={scratch.get(key)}, requires {expected}") return not reasons, reasons -def _correctness_gate( - payload: dict[str, Any], groups: int -) -> tuple[bool, list[str]]: +def _correctness_gate(payload: dict[str, Any], groups: int) -> tuple[bool, list[str]]: reasons: list[str] = [] exactness = payload.get("exactness", {}) expected_words = groups * K_M * K_D // 2 if exactness.get("word_dtype") != "uint32 packed fp16": reasons.append( - "word_dtype=" - f"{exactness.get('word_dtype')}, requires uint32 packed fp16" + f"word_dtype={exactness.get('word_dtype')}, requires uint32 packed fp16" ) if exactness.get("word_count") != expected_words: reasons.append( - "word_count=" - f"{exactness.get('word_count')}, requires {expected_words}" + f"word_count={exactness.get('word_count')}, requires {expected_words}" ) if exactness.get("full_output") is not True: reasons.append("C++ probe did not compare the full output") if exactness.get("bitwise_equal") is not True: reasons.append("candidate output is not bitwise equal to baseline") if exactness.get("mismatch_words") != 0: - reasons.append( - "mismatch_words=" - f"{exactness.get('mismatch_words')}, requires 0" - ) + reasons.append(f"mismatch_words={exactness.get('mismatch_words')}, requires 0") xor = exactness.get("xor", {}) if xor.get("max_word") != 0 or xor.get("reduction") != 0: reasons.append("packed uint32 XOR result is nonzero") return not reasons, reasons -def _performance_gate( - payload: dict[str, Any], rounds: int -) -> tuple[bool, list[str]]: +def _performance_gate(payload: dict[str, Any], rounds: int) -> tuple[bool, list[str]]: reasons: list[str] = [] pairs = payload.get("pairs", {}) if pairs.get("count") != rounds: @@ -430,13 +410,10 @@ def _performance_gate( f"candidate_faster={pairs.get('candidate_faster')}, " f"requires >= {required_wins}/{rounds}" ) - speedup_pct = payload.get("timing", {}).get( - "candidate_speedup_vs_baseline_pct" - ) + speedup_pct = payload.get("timing", {}).get("candidate_speedup_vs_baseline_pct") if speedup_pct is None or speedup_pct < K_REQUIRED_SPEEDUP_PCT: reasons.append( - f"median speedup={speedup_pct}, requires " - f">= {K_REQUIRED_SPEEDUP_PCT}%" + f"median speedup={speedup_pct}, requires >= {K_REQUIRED_SPEEDUP_PCT}%" ) return not reasons, reasons @@ -502,9 +479,7 @@ def _run_case( resource_passed, resource_reasons = _resource_gate( payload, sass, ptxas_baseline, ptxas_candidate, build_command ) - performance_passed, performance_reasons = _performance_gate( - payload, args.rounds - ) + performance_passed, performance_reasons = _performance_gate(payload, args.rounds) reasons = [ *([] if completed.returncode == 0 else ["C++ benchmark returned nonzero"]), *correctness_reasons, @@ -625,14 +600,11 @@ def _run_ncu_kernel( "l1tex__data_bank_conflicts_pipe_lsu_mem_shared.sum" ), "duration_ns": metrics.get("gpu__time_duration.sum"), - "tensor_instructions": metrics.get( - "smsp__inst_executed_pipe_tensor.sum" - ), + "tensor_instructions": metrics.get("smsp__inst_executed_pipe_tensor.sum"), }, "stdout_tail": completed.stdout[-3000:], "stderr_tail": completed.stderr[-3000:], - "permission_error": "ERR_NVGPUCTRPERM" - in (completed.stdout + completed.stderr), + "permission_error": "ERR_NVGPUCTRPERM" in (completed.stdout + completed.stderr), } @@ -659,17 +631,13 @@ def _ncu_payload( def _aggregate_gates(cases: list[dict[str, Any]]) -> dict[str, Any]: - correctness_all_cases = all( - case["gates"]["correctness_passed"] for case in cases - ) + correctness_all_cases = all(case["gates"]["correctness_passed"] for case in cases) resource_all_cases = all(case["gates"]["resource_passed"] for case in cases) all_case_performance_passed = all( case["gates"]["performance_passed"] for case in cases ) aggregate_passed = ( - correctness_all_cases - and resource_all_cases - and all_case_performance_passed + correctness_all_cases and resource_all_cases and all_case_performance_passed ) reasons: list[str] = [] if not correctness_all_cases: @@ -735,8 +703,7 @@ def main() -> int: "benchmark": "sm70_native_bm32_allp_scratch_micro", "target": "sm_70", "matrix": [ - {"pattern": pattern, "nblocks": nblocks} - for pattern, nblocks in K_CASES + {"pattern": pattern, "nblocks": nblocks} for pattern, nblocks in K_CASES ], "configuration": { "groups": args.groups, diff --git a/benchmarks/benchmark_sm70_native_bm32_full_phase_allp_micro.py b/benchmarks/benchmark_sm70_native_bm32_full_phase_allp_micro.py index 0a222682e1..cbbe4b867d 100644 --- a/benchmarks/benchmark_sm70_native_bm32_full_phase_allp_micro.py +++ b/benchmarks/benchmark_sm70_native_bm32_full_phase_allp_micro.py @@ -87,10 +87,7 @@ def _require_environment(args: argparse.Namespace) -> None: if args.rounds < K_DEFAULT_ROUNDS: raise ValueError(f"rounds must be at least {K_DEFAULT_ROUNDS}.") if args.launches_per_sample < K_DEFAULT_LAUNCHES: - raise ValueError( - "launches-per-sample must be at least " - f"{K_DEFAULT_LAUNCHES}." - ) + raise ValueError(f"launches-per-sample must be at least {K_DEFAULT_LAUNCHES}.") def _graphics_clock_mhz(physical_gpu: int) -> int: @@ -239,8 +236,7 @@ def _inspect_sass(binary: Path) -> dict[str, Any]: "ctas_per_group": ctas_per_group, "instructions": instructions, "instructions_per_group": { - name: count * ctas_per_group - for name, count in instructions.items() + name: count * ctas_per_group for name, count in instructions.items() }, } return {"available": True, "binary": str(binary), "kernels": kernels} @@ -266,8 +262,7 @@ def _resource_gate( reasons.append("C++ runtime resource gate did not pass") if baseline.get("threads_per_cta") != 512: reasons.append( - "baseline threads_per_cta=" - f"{baseline.get('threads_per_cta')}, requires 512" + f"baseline threads_per_cta={baseline.get('threads_per_cta')}, requires 512" ) if baseline.get("active_ctas_per_sm", 0) < 2: reasons.append( @@ -281,18 +276,15 @@ def _resource_gate( ) if candidate.get("registers_per_thread", 65) > 64: reasons.append( - "candidate REG=" - f"{candidate.get('registers_per_thread')}, requires <=64" + f"candidate REG={candidate.get('registers_per_thread')}, requires <=64" ) if candidate.get("local_bytes_per_thread") != 0: reasons.append( - "candidate LOCAL=" - f"{candidate.get('local_bytes_per_thread')}, requires 0" + f"candidate LOCAL={candidate.get('local_bytes_per_thread')}, requires 0" ) if candidate.get("static_shared_bytes", 48 * 1024 + 1) > 48 * 1024: reasons.append( - "candidate shared=" - f"{candidate.get('static_shared_bytes')}, requires <=49152" + f"candidate shared={candidate.get('static_shared_bytes')}, requires <=49152" ) if candidate.get("active_ctas_per_sm", 0) < 2: reasons.append( @@ -348,40 +340,31 @@ def _resource_gate( return not reasons, reasons -def _correctness_gate( - payload: dict[str, Any], groups: int -) -> tuple[bool, list[str]]: +def _correctness_gate(payload: dict[str, Any], groups: int) -> tuple[bool, list[str]]: reasons: list[str] = [] exactness = payload.get("exactness", {}) expected_words = groups * K_M * K_D // 2 if exactness.get("word_dtype") != "uint32 packed fp16": reasons.append( - "word_dtype=" - f"{exactness.get('word_dtype')}, requires uint32 packed fp16" + f"word_dtype={exactness.get('word_dtype')}, requires uint32 packed fp16" ) if exactness.get("word_count") != expected_words: reasons.append( - "word_count=" - f"{exactness.get('word_count')}, requires {expected_words}" + f"word_count={exactness.get('word_count')}, requires {expected_words}" ) if exactness.get("full_output") is not True: reasons.append("C++ probe did not compare the full output") if exactness.get("bitwise_equal") is not True: reasons.append("candidate output is not bitwise equal to baseline") if exactness.get("mismatch_words") != 0: - reasons.append( - "mismatch_words=" - f"{exactness.get('mismatch_words')}, requires 0" - ) + reasons.append(f"mismatch_words={exactness.get('mismatch_words')}, requires 0") xor = exactness.get("xor", {}) if xor.get("max_word") != 0 or xor.get("reduction") != 0: reasons.append("packed uint32 XOR result is nonzero") return not reasons, reasons -def _performance_gate( - payload: dict[str, Any], rounds: int -) -> tuple[bool, list[str]]: +def _performance_gate(payload: dict[str, Any], rounds: int) -> tuple[bool, list[str]]: reasons: list[str] = [] timing = payload.get("timing", {}) pairs = payload.get("pairs", {}) @@ -390,8 +373,7 @@ def _performance_gate( reasons.append("candidate speedup was not reported") elif speedup < K_REQUIRED_SPEEDUP_PCT: reasons.append( - f"candidate speedup {speedup:.3f}% is below " - f"{K_REQUIRED_SPEEDUP_PCT:.3f}%" + f"candidate speedup {speedup:.3f}% is below {K_REQUIRED_SPEEDUP_PCT:.3f}%" ) if pairs.get("count") != rounds: reasons.append(f"pair count={pairs.get('count')}, requires {rounds}") @@ -468,9 +450,7 @@ def _run_case( resource_passed, resource_reasons = _resource_gate( payload, sass, ptxas_candidate, build_command ) - performance_passed, performance_reasons = _performance_gate( - payload, args.rounds - ) + performance_passed, performance_reasons = _performance_gate(payload, args.rounds) reasons = [ *([] if completed.returncode == 0 else ["C++ benchmark returned nonzero"]), *correctness_reasons, @@ -584,14 +564,11 @@ def _run_ncu_kernel( "long_scoreboard_per_warp_active": metrics.get( "smsp__warp_issue_stalled_long_scoreboard_per_warp_active" ), - "tensor_instructions": metrics.get( - "smsp__inst_executed_pipe_tensor.sum" - ), + "tensor_instructions": metrics.get("smsp__inst_executed_pipe_tensor.sum"), }, "stdout_tail": completed.stdout[-3000:], "stderr_tail": completed.stderr[-3000:], - "permission_error": "ERR_NVGPUCTRPERM" - in (completed.stdout + completed.stderr), + "permission_error": "ERR_NVGPUCTRPERM" in (completed.stdout + completed.stderr), } @@ -620,9 +597,7 @@ def _ncu_payload( def _aggregate_gates(cases: list[dict[str, Any]]) -> dict[str, Any]: primary_cases = [case for case in cases if case["nblocks"] == 4] short_cases = [case for case in cases if case["nblocks"] != 4] - correctness_all_cases = all( - case["gates"]["correctness_passed"] for case in cases - ) + correctness_all_cases = all(case["gates"]["correctness_passed"] for case in cases) resource_all_cases = all(case["gates"]["resource_passed"] for case in cases) primary_performance_passed = all( case["gates"]["performance_passed"] for case in primary_cases @@ -631,9 +606,7 @@ def _aggregate_gates(cases: list[dict[str, Any]]) -> dict[str, Any]: case["gates"]["performance_passed"] for case in cases ) short_case_performance_failures = [ - case["name"] - for case in short_cases - if not case["gates"]["performance_passed"] + case["name"] for case in short_cases if not case["gates"]["performance_passed"] ] aggregate_passed = ( correctness_all_cases and resource_all_cases and primary_performance_passed @@ -702,8 +675,7 @@ def main() -> int: "benchmark": "sm70_native_bm32_full_phase_allp_micro", "target": "sm_70", "matrix": [ - {"pattern": pattern, "nblocks": nblocks} - for pattern, nblocks in K_CASES + {"pattern": pattern, "nblocks": nblocks} for pattern, nblocks in K_CASES ], "configuration": { "groups": args.groups, diff --git a/benchmarks/benchmark_sm70_native_bm32_full_phase_micro.py b/benchmarks/benchmark_sm70_native_bm32_full_phase_micro.py index fa97a71a36..d856e33db1 100644 --- a/benchmarks/benchmark_sm70_native_bm32_full_phase_micro.py +++ b/benchmarks/benchmark_sm70_native_bm32_full_phase_micro.py @@ -87,10 +87,7 @@ def _require_environment(args: argparse.Namespace) -> None: if args.rounds < K_DEFAULT_ROUNDS: raise ValueError(f"rounds must be at least {K_DEFAULT_ROUNDS}.") if args.launches_per_sample < K_DEFAULT_LAUNCHES: - raise ValueError( - "launches-per-sample must be at least " - f"{K_DEFAULT_LAUNCHES}." - ) + raise ValueError(f"launches-per-sample must be at least {K_DEFAULT_LAUNCHES}.") def _graphics_clock_mhz(physical_gpu: int) -> int: @@ -239,8 +236,7 @@ def _inspect_sass(binary: Path) -> dict[str, Any]: "ctas_per_group": ctas_per_group, "instructions": instructions, "instructions_per_group": { - name: count * ctas_per_group - for name, count in instructions.items() + name: count * ctas_per_group for name, count in instructions.items() }, } return {"available": True, "binary": str(binary), "kernels": kernels} @@ -266,8 +262,7 @@ def _resource_gate( reasons.append("C++ runtime resource gate did not pass") if baseline.get("threads_per_cta") != 512: reasons.append( - "baseline threads_per_cta=" - f"{baseline.get('threads_per_cta')}, requires 512" + f"baseline threads_per_cta={baseline.get('threads_per_cta')}, requires 512" ) if baseline.get("active_ctas_per_sm", 0) < 2: reasons.append( @@ -281,18 +276,15 @@ def _resource_gate( ) if candidate.get("registers_per_thread", 65) > 64: reasons.append( - "candidate REG=" - f"{candidate.get('registers_per_thread')}, requires <=64" + f"candidate REG={candidate.get('registers_per_thread')}, requires <=64" ) if candidate.get("local_bytes_per_thread") != 0: reasons.append( - "candidate LOCAL=" - f"{candidate.get('local_bytes_per_thread')}, requires 0" + f"candidate LOCAL={candidate.get('local_bytes_per_thread')}, requires 0" ) if candidate.get("static_shared_bytes", 48 * 1024 + 1) > 48 * 1024: reasons.append( - "candidate shared=" - f"{candidate.get('static_shared_bytes')}, requires <=49152" + f"candidate shared={candidate.get('static_shared_bytes')}, requires <=49152" ) if candidate.get("active_ctas_per_sm", 0) < 2: reasons.append( @@ -348,40 +340,31 @@ def _resource_gate( return not reasons, reasons -def _correctness_gate( - payload: dict[str, Any], groups: int -) -> tuple[bool, list[str]]: +def _correctness_gate(payload: dict[str, Any], groups: int) -> tuple[bool, list[str]]: reasons: list[str] = [] exactness = payload.get("exactness", {}) expected_words = groups * K_M * K_D // 2 if exactness.get("word_dtype") != "uint32 packed fp16": reasons.append( - "word_dtype=" - f"{exactness.get('word_dtype')}, requires uint32 packed fp16" + f"word_dtype={exactness.get('word_dtype')}, requires uint32 packed fp16" ) if exactness.get("word_count") != expected_words: reasons.append( - "word_count=" - f"{exactness.get('word_count')}, requires {expected_words}" + f"word_count={exactness.get('word_count')}, requires {expected_words}" ) if exactness.get("full_output") is not True: reasons.append("C++ probe did not compare the full output") if exactness.get("bitwise_equal") is not True: reasons.append("candidate output is not bitwise equal to baseline") if exactness.get("mismatch_words") != 0: - reasons.append( - "mismatch_words=" - f"{exactness.get('mismatch_words')}, requires 0" - ) + reasons.append(f"mismatch_words={exactness.get('mismatch_words')}, requires 0") xor = exactness.get("xor", {}) if xor.get("max_word") != 0 or xor.get("reduction") != 0: reasons.append("packed uint32 XOR result is nonzero") return not reasons, reasons -def _performance_gate( - payload: dict[str, Any], rounds: int -) -> tuple[bool, list[str]]: +def _performance_gate(payload: dict[str, Any], rounds: int) -> tuple[bool, list[str]]: reasons: list[str] = [] timing = payload.get("timing", {}) pairs = payload.get("pairs", {}) @@ -390,8 +373,7 @@ def _performance_gate( reasons.append("candidate speedup was not reported") elif speedup < K_REQUIRED_SPEEDUP_PCT: reasons.append( - f"candidate speedup {speedup:.3f}% is below " - f"{K_REQUIRED_SPEEDUP_PCT:.3f}%" + f"candidate speedup {speedup:.3f}% is below {K_REQUIRED_SPEEDUP_PCT:.3f}%" ) if pairs.get("count") != rounds: reasons.append(f"pair count={pairs.get('count')}, requires {rounds}") @@ -468,9 +450,7 @@ def _run_case( resource_passed, resource_reasons = _resource_gate( payload, sass, ptxas_candidate, build_command ) - performance_passed, performance_reasons = _performance_gate( - payload, args.rounds - ) + performance_passed, performance_reasons = _performance_gate(payload, args.rounds) reasons = [ *([] if completed.returncode == 0 else ["C++ benchmark returned nonzero"]), *correctness_reasons, @@ -584,14 +564,11 @@ def _run_ncu_kernel( "long_scoreboard_per_warp_active": metrics.get( "smsp__warp_issue_stalled_long_scoreboard_per_warp_active" ), - "tensor_instructions": metrics.get( - "smsp__inst_executed_pipe_tensor.sum" - ), + "tensor_instructions": metrics.get("smsp__inst_executed_pipe_tensor.sum"), }, "stdout_tail": completed.stdout[-3000:], "stderr_tail": completed.stderr[-3000:], - "permission_error": "ERR_NVGPUCTRPERM" - in (completed.stdout + completed.stderr), + "permission_error": "ERR_NVGPUCTRPERM" in (completed.stdout + completed.stderr), } @@ -620,9 +597,7 @@ def _ncu_payload( def _aggregate_gates(cases: list[dict[str, Any]]) -> dict[str, Any]: primary_cases = [case for case in cases if case["nblocks"] == 4] short_cases = [case for case in cases if case["nblocks"] != 4] - correctness_all_cases = all( - case["gates"]["correctness_passed"] for case in cases - ) + correctness_all_cases = all(case["gates"]["correctness_passed"] for case in cases) resource_all_cases = all(case["gates"]["resource_passed"] for case in cases) primary_performance_passed = all( case["gates"]["performance_passed"] for case in primary_cases @@ -631,9 +606,7 @@ def _aggregate_gates(cases: list[dict[str, Any]]) -> dict[str, Any]: case["gates"]["performance_passed"] for case in cases ) short_case_performance_failures = [ - case["name"] - for case in short_cases - if not case["gates"]["performance_passed"] + case["name"] for case in short_cases if not case["gates"]["performance_passed"] ] aggregate_passed = ( correctness_all_cases and resource_all_cases and primary_performance_passed @@ -702,8 +675,7 @@ def main() -> int: "benchmark": "sm70_native_bm32_full_phase_micro", "target": "sm_70", "matrix": [ - {"pattern": pattern, "nblocks": nblocks} - for pattern, nblocks in K_CASES + {"pattern": pattern, "nblocks": nblocks} for pattern, nblocks in K_CASES ], "configuration": { "groups": args.groups, diff --git a/benchmarks/benchmark_sm70_native_bm32_pv_panel_micro.py b/benchmarks/benchmark_sm70_native_bm32_pv_panel_micro.py index 3394ea1fd6..0391ec6b96 100644 --- a/benchmarks/benchmark_sm70_native_bm32_pv_panel_micro.py +++ b/benchmarks/benchmark_sm70_native_bm32_pv_panel_micro.py @@ -108,9 +108,7 @@ def _build_binary(verbose: bool) -> tuple[Path, list[str], str]: if nvcc is None: raise RuntimeError("nvcc is required to build the SM70 BM32 PV microbenchmark.") source = ( - Path(__file__).resolve().parent - / "csrc" - / "sm70_native_bm32_pv_panel_micro.cu" + Path(__file__).resolve().parent / "csrc" / "sm70_native_bm32_pv_panel_micro.cu" ) build_dir = ( Path(tempfile.gettempdir()) @@ -228,8 +226,7 @@ def _resource_gate( ) if baseline.get("threads_per_cta") != 512: reasons.append( - "baseline threads_per_cta=" - f"{baseline.get('threads_per_cta')}, requires 512" + f"baseline threads_per_cta={baseline.get('threads_per_cta')}, requires 512" ) if baseline.get("resident_active_warps") != 32: reasons.append( @@ -258,13 +255,11 @@ def _resource_gate( ) if candidate.get("registers_per_thread", 65) > 64: reasons.append( - "candidate REG=" - f"{candidate.get('registers_per_thread')}, requires <=64" + f"candidate REG={candidate.get('registers_per_thread')}, requires <=64" ) if candidate.get("local_bytes_per_thread") != 0: reasons.append( - "candidate LOCAL=" - f"{candidate.get('local_bytes_per_thread')}, requires 0" + f"candidate LOCAL={candidate.get('local_bytes_per_thread')}, requires 0" ) if candidate_topology.get("ctas_per_group") != 1: reasons.append( @@ -388,9 +383,7 @@ def _run_ncu_kernel( "long_scoreboard_per_warp_active": metrics.get( "smsp__warp_issue_stalled_long_scoreboard_per_warp_active" ), - "tensor_instructions": metrics.get( - "smsp__inst_executed_pipe_tensor.sum" - ), + "tensor_instructions": metrics.get("smsp__inst_executed_pipe_tensor.sum"), }, "stdout_tail": result.stdout[-3000:], "stderr_tail": result.stderr[-3000:], diff --git a/benchmarks/benchmark_sm70_nvfp4_gemm_micro.py b/benchmarks/benchmark_sm70_nvfp4_gemm_micro.py index e3751f9ee6..253ba1d1f9 100644 --- a/benchmarks/benchmark_sm70_nvfp4_gemm_micro.py +++ b/benchmarks/benchmark_sm70_nvfp4_gemm_micro.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Microbenchmark SM70 TurboMind NVFP4 dense GEMMs for 27B decode. +"""Microbenchmark SM70 TurboMind and QPN4 NVFP4 dense GEMMs. The default case set models one Qwen3.5-27B-NVFP4 decode token on one TP rank with tensor_parallel_size=2. It times only nvfp4_gemm_sm70_out after synthetic weights have been prepared and dispatch has been warmed, so the weighted total -is comparable to the Nsight "TurboMind NVFP4 GEMM" critical-path bucket. +is comparable to the Nsight "TurboMind NVFP4 GEMM" critical-path bucket. The +``qpn4-prefill`` mode measures the existing bounded-workspace large-M route, +including weight dequantization on every call. """ from __future__ import annotations @@ -160,23 +162,30 @@ def _run_case( use_cuda_graph: bool, mode: str, gemv_split_k: int, + gated_silu: bool, ) -> dict[str, Any]: if case.k % group_size != 0: raise ValueError(f"{case.label}: K={case.k} not divisible by {group_size}.") qweight = _make_qweight(case.k, case.n, device) scales = _make_scales(case.k, case.n, group_size, device) x = _make_input(case.m, case.k, device) - out = torch.empty((case.m, case.n), dtype=torch.float16, device=device) + if gated_silu and case.n % 2 != 0: + raise ValueError(f"{case.label}: gated-SiLU requires even N, got {case.n}.") + output_size = case.n // 2 if gated_silu else case.n + out = torch.empty((case.m, output_size), dtype=torch.float16, device=device) k_ld = 0 q_ld = 0 tm_weight = None tm_scales = None + qpn4_codes = None + qpn4_scales = None + dense_weight = None qweight_packed = None partials = None if mode == "gemm": tm_weight, tm_scales, meta = ops.nvfp4_sm70_prepare( - qweight, scales, group_size, False + qweight, scales, group_size, gated_silu ) k_ld = int(meta[0].item()) q_ld = int(meta[1].item()) @@ -190,7 +199,21 @@ def _run_case( group_size, k_ld, q_ld, + gated_silu, + ) + elif mode == "qpn4-prefill": + qpn4_codes, qpn4_scales = ops.nvfp4_qpn4_prepare_sm70(qweight, scales) + dense_weight = torch.empty((case.k, case.n), dtype=torch.float16, device=device) + run = partial( + ops.nvfp4_qpn4_dispatch_sm70_out, + out, + dense_weight.data_ptr(), + x, + qpn4_codes, + qpn4_scales, + 0.0, False, + gated_silu, ) elif mode in ("raw-gemv", "raw-gemv-warp", "raw-gemv-h2"): if case.m != 1: @@ -254,6 +277,7 @@ def _run_case( "n": case.n, "k": case.k, "group_size": group_size, + "gated_silu": gated_silu, "gemv_split_k": gemv_split_k if mode in ("raw-gemv", "raw-gemv-h2") else 0, "k_ld": k_ld, "q_ld": q_ld, @@ -266,15 +290,31 @@ def _run_case( else ( f"raw_gemv_h2_split{gemv_split_k}_{case.m}x{case.n}x{case.k}" if mode == "raw-gemv-h2" - else f"sm70_f16_nvfp4k{group_size}_f16_tnt_fff_" - f"{case.m}x{case.n}x{case.k}_1" + else ( + f"qpn4_prefill_{case.m}x{case.n}x{case.k}" + if mode == "qpn4-prefill" + else f"sm70_f16_nvfp4k{group_size}_f16_tnt_fff_" + f"{case.m}x{case.n}x{case.k}_1" + ) ) ) ), **timing, "weighted_mean_ms": weighted_mean_ms, } - del qweight, scales, tm_weight, tm_scales, qweight_packed, partials, x, out + del ( + qweight, + scales, + tm_weight, + tm_scales, + qpn4_codes, + qpn4_scales, + dense_weight, + qweight_packed, + partials, + x, + out, + ) torch.accelerator.empty_cache() return row @@ -289,6 +329,7 @@ def _write_csv(path: Path, rows: list[dict[str, Any]]) -> None: "n", "k", "group_size", + "gated_silu", "gemv_split_k", "k_ld", "q_ld", @@ -321,10 +362,21 @@ def _parse_args() -> argparse.Namespace: ) parser.add_argument( "--mode", - choices=("gemm", "raw-gemv", "raw-gemv-warp", "raw-gemv-h2"), + choices=( + "gemm", + "qpn4-prefill", + "raw-gemv", + "raw-gemv-warp", + "raw-gemv-h2", + ), default="gemm", help="Operator implementation to benchmark.", ) + parser.add_argument( + "--gated-silu", + action="store_true", + help="Fuse gate/up SiLU and emit N/2 output columns.", + ) parser.add_argument( "--gemv-split-k", type=int, @@ -351,6 +403,16 @@ def main() -> int: raise RuntimeError("Missing _C::nvfp4_sm70_prepare.") if not hasattr(torch.ops._C, "nvfp4_gemm_sm70_out"): raise RuntimeError("Missing _C::nvfp4_gemm_sm70_out.") + if args.mode == "qpn4-prefill": + required_qpn4_ops = ( + "nvfp4_qpn4_prepare_sm70", + "nvfp4_qpn4_dispatch_sm70_out", + ) + missing_qpn4_ops = [ + name for name in required_qpn4_ops if not hasattr(torch.ops._C, name) + ] + if missing_qpn4_ops: + raise RuntimeError(f"Missing QPN4 operators: {missing_qpn4_ops}.") if args.mode == "raw-gemv" and not hasattr(torch.ops._C, "nvfp4_gemv_sm70_raw_out"): raise RuntimeError("Missing _C::nvfp4_gemv_sm70_raw_out.") if args.mode == "raw-gemv-warp" and not hasattr( @@ -378,6 +440,7 @@ def main() -> int: use_cuda_graph=args.cuda_graph, mode=args.mode, gemv_split_k=args.gemv_split_k, + gated_silu=args.gated_silu, ) for case in cases ] @@ -392,6 +455,7 @@ def main() -> int: "cuda_version": torch.version.cuda, "group_size": args.group_size, "mode": args.mode, + "gated_silu": args.gated_silu, "gemv_split_k": ( args.gemv_split_k if args.mode in ("raw-gemv", "raw-gemv-h2") else 0 ), diff --git a/benchmarks/benchmark_sm70_raw_hmma_probe.py b/benchmarks/benchmark_sm70_raw_hmma_probe.py index a1928d4436..aec857c44d 100644 --- a/benchmarks/benchmark_sm70_raw_hmma_probe.py +++ b/benchmarks/benchmark_sm70_raw_hmma_probe.py @@ -32,10 +32,7 @@ def _build_binary(source: Path, verbose: bool) -> tuple[Path, list[str]]: if nvcc is None: raise RuntimeError("nvcc is required to build the SM70 raw HMMA probe.") - build_dir = ( - Path(tempfile.gettempdir()) - / f"vllm-sm70-raw-hmma-probe-{os.getuid()}" - ) + build_dir = Path(tempfile.gettempdir()) / f"vllm-sm70-raw-hmma-probe-{os.getuid()}" build_dir.mkdir(parents=True, exist_ok=True) binary = build_dir / "sm70_raw_hmma_probe_sm70" command = [ diff --git a/benchmarks/benchmark_sm70_rejection_sampler_micro.py b/benchmarks/benchmark_sm70_rejection_sampler_micro.py index 037b70653a..ec680a82dc 100644 --- a/benchmarks/benchmark_sm70_rejection_sampler_micro.py +++ b/benchmarks/benchmark_sm70_rejection_sampler_micro.py @@ -174,9 +174,7 @@ def _sparse_target_topk_topp_rejection_prototype( target_logits = sampled_logits[:num_tokens].float() bonus_logits = sampled_logits[num_tokens:].float() - target_topk_values, target_topk_ids = torch.topk( - target_logits, k=top_k, dim=-1 - ) + target_topk_values, target_topk_ids = torch.topk(target_logits, k=top_k, dim=-1) target_topk_probs = target_topk_values.softmax(dim=-1, dtype=torch.float32) target_remove_mask = target_topk_probs.cumsum(dim=-1) > top_p target_remove_mask[:, 1:] = target_remove_mask[:, :-1].clone() @@ -191,9 +189,7 @@ def _sparse_target_topk_topp_rejection_prototype( bonus_remove_mask = bonus_topk_probs.cumsum(dim=-1) > top_p bonus_remove_mask[:, 1:] = bonus_remove_mask[:, :-1].clone() bonus_remove_mask[:, 0] = False - bonus_topk_values = bonus_topk_values.masked_fill( - bonus_remove_mask, -float("inf") - ) + bonus_topk_values = bonus_topk_values.masked_fill(bonus_remove_mask, -float("inf")) bonus_topk_probs = bonus_topk_values.softmax(dim=-1, dtype=torch.float32) draft_ids_i64 = draft_token_ids.to(torch.int64).view(-1, 1) @@ -210,9 +206,9 @@ def _sparse_target_topk_topp_rejection_prototype( ).clamp_max(1.0) accepted_prefix = (torch.rand_like(accept_ratio) <= accept_ratio).cumprod(0) - residual = ( - target_topk_probs - draft_probs.gather(1, target_topk_ids) - ).clamp_min_(0.0) + residual = (target_topk_probs - draft_probs.gather(1, target_topk_ids)).clamp_min_( + 0.0 + ) residual_q = torch.empty_like(residual).exponential_() recovered_offsets = (residual / residual_q).argmax(dim=-1, keepdim=True) recovered_token_ids = target_topk_ids.gather(1, recovered_offsets).view(-1) @@ -247,9 +243,9 @@ def _full_vocab_topk_topp_probability_reference( k=min(top_k, probs.shape[-1]), dim=-1, ) - residual = ( - target_top_probs - draft_probs.gather(1, target_top_ids) - ).clamp_min_(0.0) + residual = (target_top_probs - draft_probs.gather(1, target_top_ids)).clamp_min_( + 0.0 + ) return target_draft_probs, target_top_ids, residual @@ -292,9 +288,7 @@ def _compact_tp2_topk_topp_probability_prototype( target_probs, torch.zeros_like(target_probs), ).sum(dim=-1) - residual = ( - target_probs - draft_probs.gather(1, target_ids) - ).clamp_min_(0.0) + residual = (target_probs - draft_probs.gather(1, target_ids)).clamp_min_(0.0) return target_draft_probs, target_ids, residual @@ -394,9 +388,7 @@ def main() -> None: bonus_logits_indices=torch.tensor( [num_tokens], dtype=torch.int32, device=device ), - logits_indices=torch.arange( - num_tokens + 1, dtype=torch.int32, device=device - ), + logits_indices=torch.arange(num_tokens + 1, dtype=torch.int32, device=device), ) target_logits_constrained = apply_sampling_constraints( @@ -495,9 +487,9 @@ def _set_combined_bonus(enabled: bool) -> None: args.iters, ), "target_prepare_with_clone": _time_cuda( - lambda: target_prepare_base[target_prepare_indices].to( - torch.float32 - ).clone(), + lambda: target_prepare_base[target_prepare_indices] + .to(torch.float32) + .clone(), args.warmup, args.iters, ), diff --git a/benchmarks/benchmark_sm70_source_inventory.py b/benchmarks/benchmark_sm70_source_inventory.py index f8b13b47fa..87c7e5f151 100644 --- a/benchmarks/benchmark_sm70_source_inventory.py +++ b/benchmarks/benchmark_sm70_source_inventory.py @@ -502,19 +502,14 @@ def _scan_tree(root: Path) -> dict[str, Any]: "file_count": file_count, "scanned_bytes": scanned_bytes, "env_vars": { - name: sorted(files) - for name, files in sorted(env_to_files.items()) + name: sorted(files) for name, files in sorted(env_to_files.items()) }, "runtime_env_vars": { - name: sorted(files) - for name, files in sorted(runtime_env_to_files.items()) + name: sorted(files) for name, files in sorted(runtime_env_to_files.items()) }, "env_file_counts": dict(sorted(env_counts.items())), "keyword_files": dict(sorted(keyword_files.items())), - "torch_ops": { - name: sorted(files) - for name, files in sorted(torch_ops.items()) - }, + "torch_ops": {name: sorted(files) for name, files in sorted(torch_ops.items())}, } @@ -597,11 +592,11 @@ def _compact_diff(old: dict[str, Any], latest: dict[str, Any]) -> dict[str, Any] diff["old_only_torch_op_classification"] = _classify_old_only_torch_ops( diff["torch_ops"]["old_only"] ) - diff["old_only_runtime_env_classification"] = ( - _classify_old_only_runtime_envs(diff["runtime_env_vars"]["old_only"]) + diff["old_only_runtime_env_classification"] = _classify_old_only_runtime_envs( + diff["runtime_env_vars"]["old_only"] ) - diff["old_only_keyword_file_classification"] = ( - _classify_old_only_keyword_files(diff["keyword_files"]["old_only"]) + diff["old_only_keyword_file_classification"] = _classify_old_only_keyword_files( + diff["keyword_files"]["old_only"] ) return diff @@ -617,15 +612,9 @@ def _classification_counts(section: dict[str, Any]) -> dict[str, Any]: def _inventory_gate(diff: dict[str, Any]) -> dict[str, Any]: - runtime_env = _classification_counts( - diff["old_only_runtime_env_classification"] - ) - keyword_files = _classification_counts( - diff["old_only_keyword_file_classification"] - ) - torch_ops = _classification_counts( - diff["old_only_torch_op_classification"] - ) + runtime_env = _classification_counts(diff["old_only_runtime_env_classification"]) + keyword_files = _classification_counts(diff["old_only_keyword_file_classification"]) + torch_ops = _classification_counts(diff["old_only_torch_op_classification"]) old_only_env_vars = diff["env_vars"]["old_only"] complete = ( not old_only_env_vars @@ -675,32 +664,37 @@ def main() -> int: json.dumps(payload, indent=args.indent, sort_keys=True) + "\n", encoding="utf-8", ) - print(json.dumps({ - "out": str(args.out), - "old_files": payload["old"]["file_count"], - "latest_files": payload["latest"]["file_count"], - "old_env_vars": len(payload["old"]["env_vars"]), - "latest_env_vars": len(payload["latest"]["env_vars"]), - "old_runtime_env_vars": len(payload["old"]["runtime_env_vars"]), - "latest_runtime_env_vars": len(payload["latest"]["runtime_env_vars"]), - "old_keyword_files": len(payload["old"]["keyword_files"]), - "latest_keyword_files": len(payload["latest"]["keyword_files"]), - "old_torch_ops": len(payload["old"]["torch_ops"]), - "latest_torch_ops": len(payload["latest"]["torch_ops"]), - "inventory_complete": payload["inventory_gate"]["complete"], - "old_only_env_vars": payload["inventory_gate"][ - "old_only_env_vars_count" - ], - "old_only_runtime_env_unclassified": payload["inventory_gate"][ - "old_only_runtime_env" - ]["unclassified_count"], - "old_only_keyword_files_unclassified": payload["inventory_gate"][ - "old_only_keyword_files" - ]["unclassified_count"], - "old_only_torch_ops_unclassified": payload["inventory_gate"][ - "old_only_torch_ops" - ]["unclassified_count"], - }, sort_keys=True)) + print( + json.dumps( + { + "out": str(args.out), + "old_files": payload["old"]["file_count"], + "latest_files": payload["latest"]["file_count"], + "old_env_vars": len(payload["old"]["env_vars"]), + "latest_env_vars": len(payload["latest"]["env_vars"]), + "old_runtime_env_vars": len(payload["old"]["runtime_env_vars"]), + "latest_runtime_env_vars": len(payload["latest"]["runtime_env_vars"]), + "old_keyword_files": len(payload["old"]["keyword_files"]), + "latest_keyword_files": len(payload["latest"]["keyword_files"]), + "old_torch_ops": len(payload["old"]["torch_ops"]), + "latest_torch_ops": len(payload["latest"]["torch_ops"]), + "inventory_complete": payload["inventory_gate"]["complete"], + "old_only_env_vars": payload["inventory_gate"][ + "old_only_env_vars_count" + ], + "old_only_runtime_env_unclassified": payload["inventory_gate"][ + "old_only_runtime_env" + ]["unclassified_count"], + "old_only_keyword_files_unclassified": payload["inventory_gate"][ + "old_only_keyword_files" + ]["unclassified_count"], + "old_only_torch_ops_unclassified": payload["inventory_gate"][ + "old_only_torch_ops" + ]["unclassified_count"], + }, + sort_keys=True, + ) + ) return 0 diff --git a/benchmarks/benchmark_sm70_tool_datasets.py b/benchmarks/benchmark_sm70_tool_datasets.py new file mode 100644 index 0000000000..b85077ba97 --- /dev/null +++ b/benchmarks/benchmark_sm70_tool_datasets.py @@ -0,0 +1,546 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Run fixed ToolACE and NexusRaven function-calling samples via an API. + +These are deterministic dataset subsets for matched route comparisons, not +full official leaderboard submissions. ToolACE function names are sanitized +to the OpenAI-compatible namespace while retaining their original names in +the descriptions and ground truth. +""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +import time +from pathlib import Path +from typing import Any + +import pandas as pd +import regex as re +import requests + +TOOLACE_MARKER = "Here is a list of functions in JSON format that you can invoke:" +SAFE_TOOL_NAME = re.compile(r"[^A-Za-z0-9_]") + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _post_chat(base_url: str, payload: dict[str, Any]) -> dict[str, Any]: + started = time.perf_counter() + response = requests.post( + f"{base_url.rstrip('/')}/v1/chat/completions", + json=payload, + timeout=600, + ) + if response.status_code != 200: + return { + "ok": False, + "status_code": response.status_code, + "body": response.text, + "elapsed_seconds": time.perf_counter() - started, + } + try: + body = response.json() + except requests.JSONDecodeError as exc: + return { + "ok": False, + "status_code": response.status_code, + "body": response.text, + "error": str(exc), + "elapsed_seconds": time.perf_counter() - started, + } + choice = body.get("choices", [{}])[0] + message = choice.get("message") or {} + return { + "ok": True, + "status_code": response.status_code, + "elapsed_seconds": time.perf_counter() - started, + "finish_reason": choice.get("finish_reason"), + "content": message.get("content") or "", + "reasoning_content": message.get("reasoning_content") or "", + "prompt_token_ids": body.get("prompt_token_ids") or [], + "output_token_ids": choice.get("token_ids") or [], + "tool_calls": message.get("tool_calls") or [], + "raw_response": body, + } + + +def _canonical_calls(result: dict[str, Any]) -> tuple[list[dict[str, Any]], list[str]]: + calls: list[dict[str, Any]] = [] + errors: list[str] = [] + for call in result.get("tool_calls") or []: + function = call.get("function") or {} + try: + arguments = json.loads(function.get("arguments") or "") + except (TypeError, json.JSONDecodeError) as exc: + errors.append(f"invalid arguments JSON: {exc}") + continue + calls.append({"name": function.get("name") or "", "arguments": arguments}) + return calls, errors + + +def _normalize_schema(value: Any) -> Any: + if isinstance(value, list): + return [_normalize_schema(item) for item in value] + if not isinstance(value, dict): + return value + normalized = {key: _normalize_schema(item) for key, item in value.items()} + type_map = { + "dict": "object", + "float": "number", + "tuple": "array", + "any": None, + "None": None, + } + schema_type = normalized.get("type") + if isinstance(schema_type, str) and schema_type in type_map: + mapped = type_map[schema_type] + if mapped is None: + normalized.pop("type") + else: + normalized["type"] = mapped + return normalized + + +def _stratified(items: list[Any], count: int, *, key) -> list[Any]: + items = sorted(items, key=key) + if len(items) <= count: + return items + if count == 1: + return [items[len(items) // 2]] + return [ + items[round(index * (len(items) - 1) / (count - 1))] for index in range(count) + ] + + +def _split_top_level(text: str) -> list[str]: + pieces: list[str] = [] + start = 0 + depth = 0 + quote: str | None = None + escaped = False + for index, character in enumerate(text): + if quote is not None: + if escaped: + escaped = False + elif character == "\\": + escaped = True + elif character == quote: + quote = None + elif character in {'"', "'"}: + quote = character + elif character in "([{": + depth += 1 + elif character in ")]}" and depth: + depth -= 1 + elif character == "," and depth == 0: + pieces.append(text[start:index].strip()) + start = index + 1 + pieces.append(text[start:].strip()) + return [piece for piece in pieces if piece] + + +def _toolace_functions(record: dict[str, Any]) -> list[dict[str, Any]] | None: + system = record.get("system") or "" + marker = system.find(TOOLACE_MARKER) + if marker == -1: + return None + start = system.find("[", marker + len(TOOLACE_MARKER)) + if start == -1: + return None + try: + functions = json.JSONDecoder().raw_decode(system[start:])[0] + except (TypeError, json.JSONDecodeError): + return None + if not isinstance(functions, list) or not all( + isinstance(function, dict) + and isinstance(function.get("name"), str) + and isinstance(function.get("parameters"), dict) + for function in functions + ): + return None + return functions + + +def _parse_toolace_calls( + answer: str, function_names: list[str] +) -> list[dict[str, Any]] | None: + answer = answer.strip() + if not answer.startswith("[") or not answer.endswith("]"): + return None + body = answer[1:-1].strip() + if not body: + return [] + calls: list[dict[str, Any]] = [] + for piece in _split_top_level(body): + names = [ + name + for name in function_names + if piece.startswith(f"{name}(") and piece.endswith(")") + ] + if not names: + return None + name = max(names, key=len) + argument_text = piece[len(name) + 1 : -1] + try: + call = ast.parse(f"f({argument_text})", mode="eval").body + if not isinstance(call, ast.Call) or call.args: + return None + arguments = { + keyword.arg: ast.literal_eval(keyword.value) + for keyword in call.keywords + if keyword.arg is not None + } + except (SyntaxError, ValueError): + return None + calls.append({"name": name, "arguments": arguments}) + return calls + + +def _clear_no_tool_answer(answer: str, functions: list[dict[str, Any]]) -> bool: + if not functions: + return True + lowered = answer.lower() + signatures = ( + "no available function", + "no available functions", + "lacks the required", + "lacks the parameter", + "please provide", + "provide further", + "additional information", + "not enough information", + ) + return any(signature in lowered for signature in signatures) + + +def _safe_tool_names( + functions: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], dict[str, str]]: + tools: list[dict[str, Any]] = [] + name_map: dict[str, str] = {} + used: set[str] = set() + for index, function in enumerate(functions): + slug = SAFE_TOOL_NAME.sub("_", function["name"]).strip("_") or "function" + safe_name = f"tool_{index}_{slug}"[:64] + while safe_name in used: + safe_name = f"tool_{index}_{slug}_{len(used)}"[:64] + used.add(safe_name) + name_map[function["name"]] = safe_name + tools.append( + { + "type": "function", + "function": { + "name": safe_name, + "description": ( + f"Original tool name: {function['name']}. " + f"{function.get('description') or ''}" + ), + "parameters": _normalize_schema(function["parameters"]), + }, + } + ) + return tools, name_map + + +def _run_toolace( + base_url: str, + model: str, + seed: int, + path: Path, + per_bucket: int, +) -> dict[str, Any]: + records = json.loads(path.read_text()) + buckets: dict[str, list[dict[str, Any]]] = { + "single": [], + "parallel": [], + "no_tool": [], + } + for index, record in enumerate(records): + conversations = record.get("conversations") or [] + if len(conversations) != 2 or conversations[0].get("from") != "user": + continue + functions = _toolace_functions(record) + if functions is None: + continue + answer = conversations[1].get("value") or "" + calls = _parse_toolace_calls(answer, [item["name"] for item in functions]) + if calls is not None and calls: + bucket = "parallel" if len(calls) > 1 else "single" + elif _clear_no_tool_answer(answer, functions): + bucket = "no_tool" + calls = [] + else: + continue + buckets[bucket].append( + { + "index": index, + "prompt": conversations[0]["value"], + "functions": functions, + "expected": calls, + "system_length": len(record.get("system") or ""), + } + ) + + results_by_bucket: dict[str, Any] = {} + for bucket, candidates in buckets.items(): + selected = _stratified( + candidates, + per_bucket, + key=lambda item: (item["system_length"], item["index"]), + ) + results: list[dict[str, Any]] = [] + for case in selected: + tools, name_map = _safe_tool_names(case["functions"]) + expected = [ + {"name": name_map[call["name"]], "arguments": call["arguments"]} + for call in case["expected"] + ] + payload = { + "model": model, + "messages": [{"role": "user", "content": case["prompt"]}], + "temperature": 0.0, + "seed": seed, + "max_tokens": 1024, + "return_token_ids": True, + "chat_template_kwargs": {"enable_thinking": False}, + } + if tools: + payload.update( + { + "tools": tools, + "tool_choice": "auto", + "parallel_tool_calls": True, + } + ) + result = _post_chat(base_url, payload) + actual, errors = _canonical_calls(result) + if result.get("ok") and actual != expected: + errors.append(f"calls {actual!r} != {expected!r}") + elif not result.get("ok"): + errors.append("request failed") + result.update( + { + "index": case["index"], + "expected": expected, + "actual": actual, + "errors": errors, + } + ) + results.append(result) + results_by_bucket[bucket] = { + "score": sum(not result["errors"] for result in results), + "total": len(results), + "cases": results, + } + score = sum(item["score"] for item in results_by_bucket.values()) + total = sum(item["total"] for item in results_by_bucket.values()) + return { + "score": score, + "total": total, + "protocol_passed": all( + case.get("ok") + for bucket in results_by_bucket.values() + for case in bucket["cases"] + ), + "per_bucket": results_by_bucket, + "note": "Fixed ToolACE subset with OpenAI-safe function-name mapping.", + } + + +def _nexus_tool(function: dict[str, Any]) -> dict[str, Any]: + json_types = { + "str": "string", + "int": "integer", + "list": "array", + "dict": "object", + } + properties: dict[str, Any] = {} + required: list[str] = [] + for argument in function["args_dicts"]: + properties[argument["name"]] = { + "description": argument.get("description") or "" + } + if json_type := json_types.get(str(argument.get("type"))): + properties[argument["name"]]["type"] = json_type + if argument.get("required"): + required.append(argument["name"]) + return { + "type": "function", + "function": { + "name": function["name"], + "description": function.get("description") or "", + "parameters": { + "type": "object", + "properties": properties, + "required": required, + "additionalProperties": False, + }, + }, + } + + +def _run_nexus( + base_url: str, + model: str, + seed: int, + queries_path: Path, + api_path: Path, + per_domain: int, +) -> dict[str, Any]: + queries = pd.read_parquet(queries_path) + functions = pd.read_parquet(api_path) + function_map = { + (row["dataset"], row["name"]): row for row in functions.to_dict("records") + } + domain_results: dict[str, Any] = {} + for domain in sorted(queries["dataset"].unique()): + records = queries[queries["dataset"] == domain].to_dict("records") + selected = _stratified( + list(enumerate(records)), + per_domain, + key=lambda item: (len(item[1]["prompt"]), item[0]), + ) + results: list[dict[str, Any]] = [] + for dataset_index, record in selected: + tools = [ + _nexus_tool(function_map[(domain, name)]) + for name in record["context_functions"] + ] + expected = [ + { + "name": record["python_function_name"], + "arguments": json.loads(record["python_args_dict"]), + } + ] + result = _post_chat( + base_url, + { + "model": model, + "messages": [{"role": "user", "content": record["prompt"]}], + "tools": tools, + "tool_choice": "auto", + "parallel_tool_calls": False, + "temperature": 0.0, + "seed": seed, + "max_tokens": 1024, + "return_token_ids": True, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + actual, errors = _canonical_calls(result) + if result.get("ok") and actual != expected: + errors.append(f"calls {actual!r} != {expected!r}") + elif not result.get("ok"): + errors.append("request failed") + result.update( + { + "dataset_index": dataset_index, + "prompt": record["prompt"], + "expected": expected, + "actual": actual, + "errors": errors, + } + ) + results.append(result) + domain_results[domain] = { + "score": sum(not result["errors"] for result in results), + "total": len(results), + "cases": results, + } + score = sum(item["score"] for item in domain_results.values()) + total = sum(item["total"] for item in domain_results.values()) + return { + "score": score, + "total": total, + "protocol_passed": all( + case.get("ok") + for domain in domain_results.values() + for case in domain["cases"] + ), + "per_domain": domain_results, + "note": "Fixed NexusRaven API evaluation subset; exact call/argument match.", + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", required=True) + parser.add_argument("--model", default="quality-audit-dflash2") + parser.add_argument("--seed", type=int, default=20260827) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--toolace-data", type=Path, required=True) + parser.add_argument("--toolace-per-bucket", type=int, default=4) + parser.add_argument("--nexus-queries", type=Path, required=True) + parser.add_argument("--nexus-api-list", type=Path, required=True) + parser.add_argument("--nexus-per-domain", type=int, default=4) + args = parser.parse_args() + + artifact = { + "base_url": args.base_url, + "model": args.model, + "seed": args.seed, + "sources": { + "toolace": { + "path": str(args.toolace_data), + "sha256": _sha256(args.toolace_data), + }, + "nexus_queries": { + "path": str(args.nexus_queries), + "sha256": _sha256(args.nexus_queries), + }, + "nexus_api_list": { + "path": str(args.nexus_api_list), + "sha256": _sha256(args.nexus_api_list), + }, + }, + "toolace": _run_toolace( + args.base_url, + args.model, + args.seed, + args.toolace_data, + args.toolace_per_bucket, + ), + "nexus_raven": _run_nexus( + args.base_url, + args.model, + args.seed, + args.nexus_queries, + args.nexus_api_list, + args.nexus_per_domain, + ), + } + artifact["protocol_passed"] = bool( + artifact["toolace"]["protocol_passed"] + and artifact["nexus_raven"]["protocol_passed"] + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(artifact, ensure_ascii=False, indent=2) + "\n") + print( + json.dumps( + { + "protocol_passed": artifact["protocol_passed"], + "toolace": [artifact["toolace"]["score"], artifact["toolace"]["total"]], + "nexus_raven": [ + artifact["nexus_raven"]["score"], + artifact["nexus_raven"]["total"], + ], + "out": str(args.out), + } + ) + ) + if not artifact["protocol_passed"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_sm70_tool_protocol.py b/benchmarks/benchmark_sm70_tool_protocol.py new file mode 100644 index 0000000000..b7cccb1425 --- /dev/null +++ b/benchmarks/benchmark_sm70_tool_protocol.py @@ -0,0 +1,1291 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exercise long, streamed tool calls and structured output through an API. + +This is a deterministic protocol gate, not a substitute for BFCL task scores. +It catches malformed streamed arguments, invalid tool names, history replay +failures, and JSON-schema violations while retaining the raw SSE events needed +to localize a first divergence between two server configurations. +""" + +from __future__ import annotations + +import argparse +import copy +import json +import time +from pathlib import Path +from typing import Any + +import jsonschema +import regex as re +import requests + +TOOLS: list[dict[str, Any]] = [ + { + "type": "function", + "function": { + "name": "Read", + "description": "Read a UTF-8 file.", + "parameters": { + "type": "object", + "properties": {"file_path": {"type": "string"}}, + "required": ["file_path"], + "additionalProperties": False, + }, + }, + }, + { + "type": "function", + "function": { + "name": "Grep", + "description": "Search files for a literal pattern.", + "parameters": { + "type": "object", + "properties": { + "pattern": {"type": "string"}, + "path": {"type": "string"}, + "output_mode": { + "type": "string", + "enum": ["content", "files_with_matches", "count"], + }, + }, + "required": ["pattern", "path", "output_mode"], + "additionalProperties": False, + }, + }, + }, + { + "type": "function", + "function": { + "name": "Glob", + "description": "Find files matching a glob pattern.", + "parameters": { + "type": "object", + "properties": { + "pattern": {"type": "string"}, + "path": {"type": "string"}, + }, + "required": ["pattern", "path"], + "additionalProperties": False, + }, + }, + }, + { + "type": "function", + "function": { + "name": "Bash", + "description": "Run one shell command.", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + "additionalProperties": False, + }, + }, + }, + { + "type": "function", + "function": { + "name": "Edit", + "description": "Replace an exact string in a file.", + "parameters": { + "type": "object", + "properties": { + "file_path": {"type": "string"}, + "old_string": {"type": "string"}, + "new_string": {"type": "string"}, + "replace_all": {"type": "boolean"}, + }, + "required": [ + "file_path", + "old_string", + "new_string", + "replace_all", + ], + "additionalProperties": False, + }, + }, + }, + { + "type": "function", + "function": { + "name": "Write", + "description": "Write complete UTF-8 content to a file.", + "parameters": { + "type": "object", + "properties": { + "file_path": {"type": "string"}, + "content": {"type": "string"}, + }, + "required": ["file_path", "content"], + "additionalProperties": False, + }, + }, + }, +] + + +STEPS: list[dict[str, Any]] = [ + { + "name": "Read", + "prompt": "Call Read for /tmp/project/src/app.py. Do not call another tool.", + "arguments": {"file_path": "/tmp/project/src/app.py"}, + "result": "def run():\n return 1\n", + }, + { + "name": "Grep", + "prompt": ( + "Call Grep for literal pattern run in /tmp/project/src with " + "output_mode=count. Do not call another tool." + ), + "arguments": { + "pattern": "run", + "path": "/tmp/project/src", + "output_mode": "count", + }, + "result": "/tmp/project/src/app.py:1", + }, + { + "name": "Glob", + "prompt": ( + "Call Glob with pattern **/*.py under /tmp/project. " + "Do not call another tool." + ), + "arguments": {"pattern": "**/*.py", "path": "/tmp/project"}, + "result": "/tmp/project/src/app.py", + }, + { + "name": "Bash", + "prompt": ( + "Call Bash with exactly: python -m pytest -q. Do not call another tool." + ), + "arguments": {"command": "python -m pytest -q"}, + "result": "8 passed in 1.20s", + }, + { + "name": "Edit", + "prompt": ( + "Call Edit on /tmp/project/src/app.py. Replace exactly the string " + "` return 1\\n` with ` return 2\\n`, and set replace_all=false. " + "Preserve all indentation and newlines." + ), + "arguments": { + "file_path": "/tmp/project/src/app.py", + "old_string": " return 1\n", + "new_string": " return 2\n", + "replace_all": False, + }, + "result": "updated /tmp/project/src/app.py", + }, + { + "name": "Write", + "prompt": ( + "Call Write for /tmp/project/result.json with exactly this content: " + '{"ok":true,"lines":["a\\nb","literal"]}.' + ), + "arguments": { + "file_path": "/tmp/project/result.json", + "content": '{"ok":true,"lines":["a\\nb","literal"]}', + }, + "result": "wrote 55 bytes", + }, +] + + +STRUCTURED_CASES: list[dict[str, Any]] = [ + { + "name": "nested", + "prompt": ( + "Return the build record for project alpha, successful, with stages " + "compile=12 and test=8." + ), + "schema": { + "type": "object", + "properties": { + "project": {"type": "string", "const": "alpha"}, + "success": {"type": "boolean", "const": True}, + "stages": { + "type": "array", + "prefixItems": [ + { + "type": "object", + "properties": { + "name": {"const": "compile"}, + "seconds": {"type": "integer", "const": 12}, + }, + "required": ["name", "seconds"], + "additionalProperties": False, + }, + { + "type": "object", + "properties": { + "name": {"const": "test"}, + "seconds": {"type": "integer", "const": 8}, + }, + "required": ["name", "seconds"], + "additionalProperties": False, + }, + ], + "minItems": 2, + "maxItems": 2, + }, + }, + "required": ["project", "success", "stages"], + "additionalProperties": False, + }, + }, + { + "name": "escaped_text", + "prompt": "Return path /tmp/a, the two-line text `a\\nb`, and count 2.", + "schema": { + "type": "object", + "properties": { + "path": {"type": "string", "const": "/tmp/a"}, + "text": {"type": "string", "const": "a\nb"}, + "count": {"type": "integer", "const": 2}, + }, + "required": ["path", "text", "count"], + "additionalProperties": False, + }, + }, +] + + +STREAM_PARITY_CASES: list[dict[str, str]] = [ + { + "name": "quoted_ascii", + "expected": 'The "implied" volatility is a "forward-looking" measure.', + }, + { + "name": "json_utf8", + "expected": '{"path":"/tmp/中.txt","text":"a\\nb","emoji":"🌶️"}', + }, +] + + +def _schema_by_tool_name() -> dict[str, dict[str, Any]]: + return {item["function"]["name"]: item["function"]["parameters"] for item in TOOLS} + + +def _stream_chat(base_url: str, payload: dict[str, Any]) -> dict[str, Any]: + started = time.perf_counter() + response = requests.post( + f"{base_url.rstrip('/')}/v1/chat/completions", + json=payload, + stream=True, + timeout=600, + ) + if response.status_code != 200: + return { + "ok": False, + "status_code": response.status_code, + "body": response.text, + "elapsed_seconds": time.perf_counter() - started, + } + + raw_events: list[dict[str, Any]] = [] + prompt_token_ids: list[int] = [] + output_token_ids: list[int] = [] + output_token_chunks: list[list[int]] = [] + content = "" + reasoning = "" + finish_reason = None + calls: dict[int, dict[str, Any]] = {} + malformed_sse: list[str] = [] + for raw_line in response.iter_lines(decode_unicode=True): + if not raw_line or not raw_line.startswith("data:"): + continue + data = raw_line[5:].strip() + if data == "[DONE]": + break + try: + event = json.loads(data) + except json.JSONDecodeError: + malformed_sse.append(data) + continue + raw_events.append(event) + choice = event.get("choices", [{}])[0] + if event.get("prompt_token_ids"): + prompt_token_ids.extend(event["prompt_token_ids"]) + token_ids = choice.get("token_ids") or [] + if token_ids: + output_token_chunks.append(token_ids) + output_token_ids.extend(token_ids) + finish_reason = choice.get("finish_reason") or finish_reason + delta = choice.get("delta") or {} + content += delta.get("content") or "" + reasoning += delta.get("reasoning_content") or "" + for call_delta in delta.get("tool_calls") or []: + index = int(call_delta.get("index", 0)) + call = calls.setdefault( + index, + { + "id": "", + "type": "function", + "function": {"name": "", "arguments": ""}, + }, + ) + call["id"] += call_delta.get("id") or "" + function = call_delta.get("function") or {} + call["function"]["name"] += function.get("name") or "" + call["function"]["arguments"] += function.get("arguments") or "" + + return { + "ok": not malformed_sse, + "status_code": response.status_code, + "elapsed_seconds": time.perf_counter() - started, + "finish_reason": finish_reason, + "content": content, + "reasoning_content": reasoning, + "prompt_token_ids": prompt_token_ids, + "output_token_ids": output_token_ids, + "output_token_chunks": output_token_chunks, + "tool_calls": [calls[index] for index in sorted(calls)], + "malformed_sse": malformed_sse, + "raw_events": raw_events, + } + + +def _nonstream_chat(base_url: str, payload: dict[str, Any]) -> dict[str, Any]: + started = time.perf_counter() + response = requests.post( + f"{base_url.rstrip('/')}/v1/chat/completions", + json=payload, + timeout=600, + ) + if response.status_code != 200: + return { + "ok": False, + "status_code": response.status_code, + "body": response.text, + "elapsed_seconds": time.perf_counter() - started, + } + try: + body = response.json() + except requests.JSONDecodeError as exc: + return { + "ok": False, + "status_code": response.status_code, + "body": response.text, + "error": str(exc), + "elapsed_seconds": time.perf_counter() - started, + } + choice = body.get("choices", [{}])[0] + message = choice.get("message") or {} + return { + "ok": True, + "status_code": response.status_code, + "elapsed_seconds": time.perf_counter() - started, + "finish_reason": choice.get("finish_reason"), + "content": message.get("content") or "", + "reasoning_content": message.get("reasoning_content") or "", + "prompt_token_ids": body.get("prompt_token_ids") or [], + "output_token_ids": choice.get("token_ids") or [], + "tool_calls": message.get("tool_calls") or [], + "raw_response": body, + } + + +def _responses_tools() -> list[dict[str, Any]]: + tools: list[dict[str, Any]] = [] + for tool in TOOLS: + function = tool["function"] + tools.append( + { + "type": "function", + "name": function["name"], + "description": function["description"], + "parameters": function["parameters"], + "strict": True, + } + ) + return tools + + +def _stream_responses(base_url: str, payload: dict[str, Any]) -> dict[str, Any]: + """Collect Responses SSE without relying on a particular SDK version.""" + started = time.perf_counter() + response = requests.post( + f"{base_url.rstrip('/')}/v1/responses", + json=payload, + stream=True, + timeout=600, + ) + if response.status_code != 200: + return { + "ok": False, + "status_code": response.status_code, + "body": response.text, + "elapsed_seconds": time.perf_counter() - started, + } + + raw_events: list[dict[str, Any]] = [] + malformed_sse: list[str] = [] + response_id = "" + response_status = "" + output_text_chunks: list[str] = [] + protocol_errors: list[str] = [] + calls: dict[int, dict[str, Any]] = {} + + for raw_line in response.iter_lines(decode_unicode=True): + if not raw_line or not raw_line.startswith("data:"): + continue + data = raw_line[5:].strip() + if data == "[DONE]": + break + try: + event = json.loads(data) + except json.JSONDecodeError: + malformed_sse.append(data) + continue + raw_events.append(event) + event_type = event.get("type") or "" + event_response = event.get("response") or {} + response_id = event_response.get("id") or response_id + response_status = event_response.get("status") or response_status + + if event_type == "response.output_text.delta": + output_text_chunks.append(event.get("delta") or "") + elif event_type == "response.output_item.added": + item = event.get("item") or {} + if item.get("type") == "function_call": + output_index = int(event.get("output_index", 0)) + calls[output_index] = { + "output_index": output_index, + "added_item": item, + "argument_chunks": [], + "done_arguments": None, + "done_item": None, + } + elif event_type == "response.function_call_arguments.delta": + output_index = int(event.get("output_index", 0)) + call = calls.get(output_index) + if call is None: + protocol_errors.append( + f"arguments delta before item at output {output_index}" + ) + else: + call["argument_chunks"].append(event.get("delta") or "") + elif event_type == "response.function_call_arguments.done": + output_index = int(event.get("output_index", 0)) + call = calls.get(output_index) + if call is None: + protocol_errors.append( + f"arguments done before item at output {output_index}" + ) + else: + call["done_arguments"] = event.get("arguments") or "" + elif event_type == "response.output_item.done": + item = event.get("item") or {} + if item.get("type") == "function_call": + output_index = int(event.get("output_index", 0)) + call = calls.get(output_index) + if call is None: + protocol_errors.append( + f"item done before item added at output {output_index}" + ) + else: + call["done_item"] = item + elif event_type in {"response.failed", "error"}: + protocol_errors.append(json.dumps(event, ensure_ascii=False)) + + function_calls: list[dict[str, Any]] = [] + for output_index in sorted(calls): + record = calls[output_index] + added = record["added_item"] + done = record["done_item"] or {} + delta_arguments = "".join(record["argument_chunks"]) + done_arguments = record["done_arguments"] + final_arguments = done.get("arguments") + if final_arguments is None: + final_arguments = done_arguments + if final_arguments is None: + final_arguments = delta_arguments or added.get("arguments") or "" + + if done_arguments is not None and done_arguments != delta_arguments: + protocol_errors.append( + f"arguments delta/done mismatch at output {output_index}" + ) + if done and done.get("arguments") != final_arguments: + protocol_errors.append( + f"arguments done/item mismatch at output {output_index}" + ) + if done and done.get("name") != added.get("name"): + protocol_errors.append(f"tool name changed at output {output_index}") + + function_calls.append( + { + "output_index": output_index, + "id": done.get("id") or added.get("id") or "", + "call_id": done.get("call_id") or added.get("call_id") or "", + "name": done.get("name") or added.get("name") or "", + "arguments": final_arguments, + "argument_chunks": record["argument_chunks"], + "done_arguments": done_arguments, + "added_item": added, + "done_item": record["done_item"], + } + ) + + completed = any(event.get("type") == "response.completed" for event in raw_events) + if not completed: + protocol_errors.append("response.completed event missing") + return { + "ok": not malformed_sse and not protocol_errors, + "status_code": response.status_code, + "elapsed_seconds": time.perf_counter() - started, + "response_id": response_id, + "response_status": response_status, + "output_text": "".join(output_text_chunks), + "function_calls": function_calls, + "malformed_sse": malformed_sse, + "protocol_errors": protocol_errors, + "raw_events": raw_events, + } + + +def _validate_responses_tool_step( + result: dict[str, Any], expected: dict[str, Any] +) -> list[str]: + if not result.get("ok"): + return ["request failed"] + calls = result.get("function_calls") or [] + if len(calls) != 1: + return [f"expected one function call, got {len(calls)}"] + call = calls[0] + errors: list[str] = [] + if call["name"] != expected["name"]: + errors.append(f"tool name {call['name']!r} != {expected['name']!r}") + if not call.get("call_id"): + errors.append("call_id missing") + try: + arguments = json.loads(call["arguments"]) + except json.JSONDecodeError as exc: + return errors + [f"invalid arguments JSON: {exc}"] + schema = _schema_by_tool_name().get(call["name"]) + if schema is None: + errors.append(f"unknown tool {call['name']!r}") + else: + try: + jsonschema.validate(arguments, schema) + except jsonschema.ValidationError as exc: + errors.append(f"schema violation: {exc.message}") + if arguments != expected["arguments"]: + errors.append(f"arguments {arguments!r} != {expected['arguments']!r}") + return errors + + +def _responses_payload( + model: str, + seed: int, + input_items: str | list[dict[str, Any]], + tool_name: str, + *, + previous_response_id: str | None = None, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "model": model, + "input": input_items, + "instructions": ( + "Follow the user literally, call exactly the requested listed tool, " + "and preserve every string byte including whitespace and XML-like text." + ), + "tools": _responses_tools(), + "tool_choice": {"type": "function", "name": tool_name}, + "parallel_tool_calls": False, + "stream": True, + "store": True, + "temperature": 0.0, + "seed": seed, + "max_output_tokens": 2048, + "chat_template_kwargs": {"enable_thinking": False}, + } + if previous_response_id is not None: + payload["previous_response_id"] = previous_response_id + return payload + + +def _run_responses_chain(base_url: str, model: str, seed: int) -> dict[str, Any]: + """Cover stateful and explicit-history Responses tool protocols.""" + stateful_steps = [*STEPS, STEPS[0]] + stateful_results: list[dict[str, Any]] = [] + previous_response_id: str | None = None + previous_call_id = "" + previous_result = "" + for step_index, step in enumerate(stateful_steps): + if previous_response_id is None: + input_items: str | list[dict[str, Any]] = step["prompt"] + else: + input_items = [ + { + "type": "function_call_output", + "call_id": previous_call_id, + "output": previous_result, + }, + {"role": "user", "content": step["prompt"]}, + ] + payload = _responses_payload( + model, + seed, + input_items, + step["name"], + previous_response_id=previous_response_id, + ) + result = _stream_responses(base_url, payload) + errors = _validate_responses_tool_step(result, step) + result["step"] = step_index + result["expected_name"] = step["name"] + result["errors"] = errors + stateful_results.append(result) + if errors: + break + previous_response_id = result["response_id"] + previous_call_id = result["function_calls"][0]["call_id"] + previous_result = step["result"] + + explicit_input: list[dict[str, Any]] = [] + for step_index, step in enumerate(STEPS): + call_id = f"call_explicit_history_{step_index}" + explicit_input.extend( + [ + {"role": "user", "content": step["prompt"]}, + { + "type": "function_call", + "call_id": call_id, + "name": step["name"], + "arguments": json.dumps( + step["arguments"], ensure_ascii=False, separators=(",", ":") + ), + }, + { + "type": "function_call_output", + "call_id": call_id, + "output": step["result"], + }, + ] + ) + replay_step = STEPS[0] + explicit_input.append({"role": "user", "content": replay_step["prompt"]}) + explicit_result = _stream_responses( + base_url, + _responses_payload( + model, + seed, + explicit_input, + replay_step["name"], + ), + ) + explicit_errors = _validate_responses_tool_step(explicit_result, replay_step) + explicit_result["errors"] = explicit_errors + + return { + "passed": ( + len(stateful_results) == len(stateful_steps) + and all(not item["errors"] for item in stateful_results) + and not explicit_errors + ), + "stateful_steps": stateful_results, + "explicit_history_replay": explicit_result, + } + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + return [json.loads(line) for line in path.read_text().splitlines() if line] + + +def _normalize_bfcl_schema(value: Any) -> Any: + if isinstance(value, list): + return [_normalize_bfcl_schema(item) for item in value] + if not isinstance(value, dict): + return value + normalized = {key: _normalize_bfcl_schema(item) for key, item in value.items()} + type_map = {"dict": "object", "float": "number", "tuple": "array"} + if isinstance(normalized.get("type"), str): + normalized["type"] = type_map.get(normalized["type"], normalized["type"]) + return normalized + + +def _bfcl_tools(functions: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + { + "type": "function", + "function": { + "name": function["name"], + "description": function.get("description") or "", + "parameters": _normalize_bfcl_schema(function["parameters"]), + }, + } + for function in functions + ] + + +def _bfcl_normalized_value(value: Any) -> Any: + if isinstance(value, str): + return re.sub(r"[ ,./\-_*^]", "", value).lower().replace("'", '"') + if isinstance(value, list): + return [_bfcl_normalized_value(item) for item in value] + if isinstance(value, dict): + return { + key: _bfcl_normalized_value(item) for key, item in sorted(value.items()) + } + return value + + +def _bfcl_call_matches( + actual: dict[str, Any], + expected: dict[str, Any], + functions: dict[str, dict[str, Any]], +) -> tuple[bool, str]: + expected_name, expected_arguments = next(iter(expected.items())) + if actual.get("name") != expected_name: + return False, f"tool name {actual.get('name')!r} != {expected_name!r}" + arguments = actual.get("arguments") + if not isinstance(arguments, dict): + return False, "arguments are not a JSON object" + function = functions.get(expected_name) + if function is None: + return False, f"function description missing for {expected_name!r}" + properties = function["parameters"].get("properties") or {} + required = function["parameters"].get("required") or [] + for name in required: + if name not in arguments: + return False, f"required argument {name!r} missing" + for name, value in arguments.items(): + if name not in properties or name not in expected_arguments: + return False, f"unexpected argument {name!r}" + allowed = expected_arguments[name] + if not any( + _bfcl_normalized_value(value) == _bfcl_normalized_value(candidate) + for candidate in allowed + ): + return False, f"argument {name!r} value {value!r} not in {allowed!r}" + for name, allowed in expected_arguments.items(): + if name not in arguments and "" not in allowed: + return False, f"expected argument {name!r} missing" + return True, "" + + +def _validate_bfcl_calls( + result: dict[str, Any], + entry: dict[str, Any], + ground_truth: list[dict[str, Any]] | None, + *, + irrelevance: bool, +) -> list[str]: + if not result.get("ok"): + return ["request failed"] + canonical = _canonical_tool_calls(result.get("tool_calls") or []) + if irrelevance: + return [] if not canonical else [f"irrelevant request called {canonical!r}"] + assert ground_truth is not None + if len(canonical) != len(ground_truth): + return [ + f"wrong number of calls: got {len(canonical)}, expected {len(ground_truth)}" + ] + functions = {function["name"]: function for function in entry["function"]} + unmatched = set(range(len(canonical))) + errors: list[str] = [] + for expected in ground_truth: + candidate_errors: list[str] = [] + for index in list(unmatched): + matched, error = _bfcl_call_matches(canonical[index], expected, functions) + if matched: + unmatched.remove(index) + break + candidate_errors.append(error) + else: + errors.append(f"no call matched {expected!r}: {candidate_errors!r}") + return errors + + +def _run_bfcl_sample( + base_url: str, + model: str, + seed: int, + data_dir: Path, + per_category: int, +) -> dict[str, Any]: + categories = { + "simple_python": False, + "parallel": False, + "multiple": False, + "irrelevance": True, + } + category_results: dict[str, Any] = {} + for category, irrelevance in categories.items(): + filename = f"BFCL_v4_{category}.json" + entries = _read_jsonl(data_dir / filename)[:per_category] + ground_truth_by_id: dict[str, list[dict[str, Any]]] = {} + if not irrelevance: + ground_truth_by_id = { + item["id"]: item["ground_truth"] + for item in _read_jsonl(data_dir / "possible_answer" / filename) + } + results: list[dict[str, Any]] = [] + for entry in entries: + payload = { + "model": model, + "messages": entry["question"][0], + "tools": _bfcl_tools(entry["function"]), + "tool_choice": "auto", + "parallel_tool_calls": True, + "stream": True, + "return_token_ids": True, + "temperature": 0.0, + "seed": seed, + "max_tokens": 1024, + "chat_template_kwargs": {"enable_thinking": False}, + } + result = _stream_chat(base_url, payload) + errors = _validate_bfcl_calls( + result, + entry, + ground_truth_by_id.get(entry["id"]), + irrelevance=irrelevance, + ) + result["id"] = entry["id"] + result["errors"] = errors + results.append(result) + category_results[category] = { + "passed": sum(not item["errors"] for item in results), + "total": len(results), + "cases": results, + } + passed = sum(item["passed"] for item in category_results.values()) + total = sum(item["total"] for item in category_results.values()) + return { + "passed": passed == total, + "score": passed, + "total": total, + "per_category": category_results, + "note": "Fixed BFCL v4 subset; not a full leaderboard submission.", + } + + +def _adapt_openai_json_schema(schema: dict[str, Any]) -> dict[str, Any]: + adapted = copy.deepcopy(schema) + if "type" not in adapted: + adapted["type"] = "object" + + def visit(value: Any) -> None: + if isinstance(value, list): + for item in value: + visit(item) + return + if not isinstance(value, dict): + return + properties = value.get("properties") + if isinstance(properties, dict): + if value.get("additionalProperties", True): + value["additionalProperties"] = False + value["required"] = list(properties) + for item in value.values(): + visit(item) + + visit(adapted) + return adapted + + +def _run_json_schema_bench_sample( + base_url: str, + model: str, + seed: int, + schema_dir: Path, + limit: int, + max_schema_bytes: int, +) -> dict[str, Any]: + results: list[dict[str, Any]] = [] + candidates = sorted( + ( + path + for path in schema_dir.glob("*.json") + if path.stat().st_size <= max_schema_bytes + ), + key=lambda path: (path.stat().st_size, path.name), + ) + if len(candidates) > limit: + if limit == 1: + candidates = [candidates[len(candidates) // 2]] + else: + candidates = [ + candidates[round(index * (len(candidates) - 1) / (limit - 1))] + for index in range(limit) + ] + for path in candidates: + schema = _adapt_openai_json_schema(json.loads(path.read_text())) + payload = { + "model": model, + "messages": [ + { + "role": "system", + "content": ( + "You need to generate a JSON object that matches the " + "schema below." + ), + }, + { + "role": "user", + "content": json.dumps(schema, ensure_ascii=False), + }, + ], + "stream": True, + "return_token_ids": True, + "temperature": 0.0, + "seed": seed, + "max_tokens": 2048, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": path.stem, + "strict": True, + "schema": schema, + }, + }, + } + result = _stream_chat(base_url, payload) + errors: list[str] = [] + if not result.get("ok"): + errors.append("request failed") + else: + try: + value = json.loads(result["content"]) + jsonschema.validate(value, schema) + except (json.JSONDecodeError, jsonschema.ValidationError) as exc: + errors.append(str(exc)) + result["schema"] = path.name + result["errors"] = errors + results.append(result) + passed = sum(not item["errors"] for item in results) + return { + "passed": passed == len(results), + "score": passed, + "total": len(results), + "cases": results, + "note": ( + "Deterministic zero-shot WashingtonPost JSONSchemaBench subset; " + f"schemas are size-stratified up to {max_schema_bytes} bytes; not a " + "full benchmark score." + ), + } + + +def _canonical_tool_calls(calls: list[dict[str, Any]]) -> list[dict[str, Any]]: + canonical: list[dict[str, Any]] = [] + for call in calls: + function = call.get("function") or {} + arguments = function.get("arguments") or "" + try: + parsed_arguments: Any = json.loads(arguments) + except (TypeError, json.JSONDecodeError): + parsed_arguments = arguments + canonical.append( + { + "name": function.get("name") or "", + "arguments": parsed_arguments, + } + ) + return canonical + + +def _compare_stream_pair( + nonstream: dict[str, Any], + stream: dict[str, Any], + *, + compare_content: bool, + compare_tools: bool, +) -> list[str]: + errors: list[str] = [] + if not nonstream.get("ok") or not stream.get("ok"): + return ["stream or non-stream request failed"] + if compare_content and nonstream["content"] != stream["content"]: + errors.append("stream/non-stream content mismatch") + if nonstream["reasoning_content"] != stream["reasoning_content"]: + errors.append("stream/non-stream reasoning mismatch") + if compare_tools and _canonical_tool_calls( + nonstream["tool_calls"] + ) != _canonical_tool_calls(stream["tool_calls"]): + errors.append("stream/non-stream tool-call mismatch") + nonstream_ids = nonstream.get("output_token_ids") or [] + stream_ids = stream.get("output_token_ids") or [] + if not nonstream_ids or not stream_ids: + errors.append("output token IDs unavailable for parity check") + elif nonstream_ids != stream_ids: + errors.append("stream/non-stream output token IDs differ") + return errors + + +def _run_stream_parity(base_url: str, model: str, seed: int) -> dict[str, Any]: + results: list[dict[str, Any]] = [] + for case in STREAM_PARITY_CASES: + expected = case["expected"] + payload = { + "model": model, + "messages": [ + { + "role": "user", + "content": ( + "Reply with exactly the text after TEXT, without quotes, " + f"commentary, or a code fence.\nTEXT: {expected}" + ), + } + ], + "return_token_ids": True, + "temperature": 0.0, + "seed": seed, + "max_tokens": 256, + "chat_template_kwargs": {"enable_thinking": False}, + } + nonstream = _nonstream_chat(base_url, {**payload, "stream": False}) + stream = _stream_chat(base_url, {**payload, "stream": True}) + errors = _compare_stream_pair( + nonstream, + stream, + compare_content=True, + compare_tools=False, + ) + if nonstream.get("ok") and nonstream["content"] != expected: + errors.append("non-stream exact-copy output mismatch") + if stream.get("ok") and stream["content"] != expected: + errors.append("stream exact-copy output mismatch") + results.append( + { + "case": case["name"], + "expected": expected, + "nonstream": nonstream, + "stream": stream, + "errors": errors, + } + ) + + step = STEPS[0] + tool_payload = { + "model": model, + "messages": [{"role": "user", "content": step["prompt"]}], + "tools": TOOLS, + "tool_choice": "auto", + "return_token_ids": True, + "temperature": 0.0, + "seed": seed, + "max_tokens": 512, + "chat_template_kwargs": {"enable_thinking": False}, + } + nonstream = _nonstream_chat(base_url, {**tool_payload, "stream": False}) + stream = _stream_chat(base_url, {**tool_payload, "stream": True}) + errors = _compare_stream_pair( + nonstream, + stream, + compare_content=True, + compare_tools=True, + ) + if nonstream.get("ok"): + errors.extend(_validate_tool_step(nonstream, step)) + if stream.get("ok"): + errors.extend(_validate_tool_step(stream, step)) + results.append( + { + "case": "tool_read", + "nonstream": nonstream, + "stream": stream, + "errors": errors, + } + ) + return {"passed": all(not item["errors"] for item in results), "cases": results} + + +def _validate_tool_step(result: dict[str, Any], expected: dict[str, Any]) -> list[str]: + errors: list[str] = [] + calls = result.get("tool_calls") or [] + if len(calls) != 1: + return [f"expected one tool call, got {len(calls)}"] + call = calls[0] + name = call["function"]["name"] + if name != expected["name"]: + errors.append(f"tool name {name!r} != {expected['name']!r}") + try: + arguments = json.loads(call["function"]["arguments"]) + except json.JSONDecodeError as exc: + return errors + [f"invalid arguments JSON: {exc}"] + schema = _schema_by_tool_name().get(name) + if schema is None: + errors.append(f"unknown tool {name!r}") + else: + try: + jsonschema.validate(arguments, schema) + except jsonschema.ValidationError as exc: + errors.append(f"schema violation: {exc.message}") + if arguments != expected["arguments"]: + errors.append(f"arguments {arguments!r} != {expected['arguments']!r}") + return errors + + +def _run_tool_chain(base_url: str, model: str, seed: int) -> dict[str, Any]: + messages: list[dict[str, Any]] = [ + { + "role": "system", + "content": ( + "You are a tool protocol validation agent. Follow each user request " + "literally, call only a listed tool, and preserve string whitespace." + ), + } + ] + results: list[dict[str, Any]] = [] + for step_index, step in enumerate(STEPS): + messages.append({"role": "user", "content": step["prompt"]}) + payload = { + "model": model, + "messages": messages, + "tools": TOOLS, + "tool_choice": "auto", + "stream": True, + "return_token_ids": True, + "temperature": 0.0, + "seed": seed, + "max_tokens": 2048, + } + result = _stream_chat(base_url, payload) + errors = ( + _validate_tool_step(result, step) + if result.get("ok") + else ["request failed"] + ) + result["step"] = step_index + result["expected_name"] = step["name"] + result["errors"] = errors + results.append(result) + if errors: + break + + tool_call = result["tool_calls"][0] + history_tool_call = copy.deepcopy(tool_call) + history_tool_call["id"] = f"call_protocol_step_{step_index}" + messages.append( + { + "role": "assistant", + "content": result.get("content") or None, + "tool_calls": [history_tool_call], + } + ) + messages.append( + { + "role": "tool", + "tool_call_id": history_tool_call["id"], + "content": step["result"], + } + ) + return { + "passed": len(results) == len(STEPS) + and all(not item["errors"] for item in results), + "steps": results, + } + + +def _run_structured(base_url: str, model: str, seed: int) -> dict[str, Any]: + results: list[dict[str, Any]] = [] + for case in STRUCTURED_CASES: + payload = { + "model": model, + "messages": [{"role": "user", "content": case["prompt"]}], + "stream": True, + "return_token_ids": True, + "temperature": 0.0, + "seed": seed, + "max_tokens": 512, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": case["name"], + "strict": True, + "schema": case["schema"], + }, + }, + } + result = _stream_chat(base_url, payload) + errors: list[str] = [] + if not result.get("ok"): + errors.append("request failed") + else: + try: + value = json.loads(result["content"]) + jsonschema.validate(value, case["schema"]) + except (json.JSONDecodeError, jsonschema.ValidationError) as exc: + errors.append(str(exc)) + result["case"] = case["name"] + result["errors"] = errors + results.append(result) + return {"passed": all(not item["errors"] for item in results), "cases": results} + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", required=True) + parser.add_argument("--model", default="quality-audit-dflash2") + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--seed", type=int, default=20260827) + parser.add_argument("--bfcl-data-dir", type=Path) + parser.add_argument("--bfcl-per-category", type=int, default=0) + parser.add_argument("--json-schema-dir", type=Path) + parser.add_argument("--json-schema-limit", type=int, default=0) + parser.add_argument("--json-schema-max-bytes", type=int, default=16384) + parser.add_argument( + "--only-datasets", + action="store_true", + help="Skip synthetic protocol and Responses gates for matched dataset arms.", + ) + args = parser.parse_args() + + artifact: dict[str, Any] = { + "base_url": args.base_url, + "model": args.model, + "seed": args.seed, + "only_datasets": args.only_datasets, + } + if not args.only_datasets: + artifact.update( + { + "tool_chain": _run_tool_chain(args.base_url, args.model, args.seed), + "structured": _run_structured(args.base_url, args.model, args.seed), + "stream_parity": _run_stream_parity( + args.base_url, args.model, args.seed + ), + "responses": _run_responses_chain(args.base_url, args.model, args.seed), + } + ) + if args.bfcl_data_dir is not None and args.bfcl_per_category > 0: + artifact["bfcl_fixed_sample"] = _run_bfcl_sample( + args.base_url, + args.model, + args.seed, + args.bfcl_data_dir, + args.bfcl_per_category, + ) + if args.json_schema_dir is not None and args.json_schema_limit > 0: + artifact["json_schema_bench_fixed_sample"] = _run_json_schema_bench_sample( + args.base_url, + args.model, + args.seed, + args.json_schema_dir, + args.json_schema_limit, + args.json_schema_max_bytes, + ) + result_keys = { + "tool_chain", + "structured", + "stream_parity", + "responses", + "bfcl_fixed_sample", + "json_schema_bench_fixed_sample", + } + required_results = [value for key, value in artifact.items() if key in result_keys] + artifact["passed"] = bool( + required_results and all(value["passed"] for value in required_results) + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(artifact, ensure_ascii=False, indent=2) + "\n") + print(json.dumps({"passed": artifact["passed"], "out": str(args.out)})) + if not artifact["passed"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_sm70_turbomind_parity.py b/benchmarks/benchmark_sm70_turbomind_parity.py index ed1f27739f..bbc2cd7c75 100644 --- a/benchmarks/benchmark_sm70_turbomind_parity.py +++ b/benchmarks/benchmark_sm70_turbomind_parity.py @@ -94,8 +94,9 @@ def _load_awq( device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: weight_map = _weight_map(model_path) - with safe_open(model_path / weight_map[f"{layer}.qweight"], - framework="pt", device="cpu") as f: + with safe_open( + model_path / weight_map[f"{layer}.qweight"], framework="pt", device="cpu" + ) as f: qweight = f.get_tensor(f"{layer}.qweight").to(device).contiguous() scales = f.get_tensor(f"{layer}.scales").to(device).contiguous() qzeros = f.get_tensor(f"{layer}.qzeros").to(device).contiguous() @@ -108,8 +109,9 @@ def _load_fp8( device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: weight_map = _weight_map(model_path) - with safe_open(model_path / weight_map[f"{layer}.weight"], - framework="pt", device="cpu") as f: + with safe_open( + model_path / weight_map[f"{layer}.weight"], framework="pt", device="cpu" + ) as f: weight = f.get_tensor(f"{layer}.weight").to(device).contiguous() scales = ( f.get_tensor(f"{layer}.weight_scale_inv") @@ -158,9 +160,7 @@ def _dump_fp8( ops = _import_sm70_ops() weight, scales = _load_fp8(args.model, args.layer, device) x = _make_input(args.m, weight.shape[1], device) - tm_weight, tm_scales, meta = ops.fp8_sm70_prepare( - weight, scales, args.group_size - ) + tm_weight, tm_scales, meta = ops.fp8_sm70_prepare(weight, scales, args.group_size) out = torch.empty((args.m, weight.shape[0]), device=device, dtype=torch.float16) ops.fp8_gemm_sm70_out_meta(out, x, tm_weight, tm_scales, meta) details = { @@ -205,8 +205,13 @@ def _dump(args: argparse.Namespace) -> int: "output": out.detach().cpu(), } torch.save(payload, args.out) - print(json.dumps({k: v for k, v in payload.items() if k != "output"}, - indent=2, sort_keys=True)) + print( + json.dumps( + {k: v for k, v in payload.items() if k != "output"}, + indent=2, + sort_keys=True, + ) + ) return 0 @@ -270,9 +275,9 @@ def _bench_awq( tm_weight, tm_scales, meta = ops.awq_sm70_prepare( qweight, scales, qzeros, args.group_size ) - out = torch.empty((args.m, qweight.shape[1] * 8), - device=device, - dtype=torch.float16) + out = torch.empty( + (args.m, qweight.shape[1] * 8), device=device, dtype=torch.float16 + ) k_ld = int(meta[0].item()) q_ld = int(meta[1].item()) @@ -300,9 +305,7 @@ def _bench_fp8( ops = _import_sm70_ops() weight, scales = _load_fp8(args.model, args.layer, device) x = _make_input(args.m, weight.shape[1], device) - tm_weight, tm_scales, meta = ops.fp8_sm70_prepare( - weight, scales, args.group_size - ) + tm_weight, tm_scales, meta = ops.fp8_sm70_prepare(weight, scales, args.group_size) out = torch.empty((args.m, weight.shape[0]), device=device, dtype=torch.float16) def run() -> None: diff --git a/benchmarks/benchmark_sm70_wmma_fragment_probe.py b/benchmarks/benchmark_sm70_wmma_fragment_probe.py index 26e7c4435d..301d49ef27 100644 --- a/benchmarks/benchmark_sm70_wmma_fragment_probe.py +++ b/benchmarks/benchmark_sm70_wmma_fragment_probe.py @@ -15,11 +15,7 @@ def _load_extension() -> object: - source = ( - Path(__file__).resolve().parent - / "csrc" - / "sm70_wmma_fragment_probe.cu" - ) + source = Path(__file__).resolve().parent / "csrc" / "sm70_wmma_fragment_probe.cu" os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "7.0") return load( name="sm70_wmma_fragment_probe_v3", @@ -41,17 +37,16 @@ def _parse_args() -> argparse.Namespace: def main() -> int: args = _parse_args() device = torch.device(args.device) - if not torch.cuda.is_available() or torch.cuda.get_device_capability( - device - ) != (7, 0): + if not torch.cuda.is_available() or torch.cuda.get_device_capability(device) != ( + 7, + 0, + ): raise RuntimeError("This probe requires an SM70 CUDA device.") extension = _load_extension() tile = torch.arange(256, device=device, dtype=torch.float16).reshape(16, 16) words = extension.dump_matrix_a_fragment(tile, args.stride) - reference_words, swizzled_words = ( - extension.compare_swizzled_matrix_a_fragment(tile) - ) + reference_words, swizzled_words = extension.compare_swizzled_matrix_a_fragment(tile) matrix_b_row_words = extension.dump_matrix_b_fragment(tile, False) matrix_b_col_storage = tile.t().contiguous() matrix_b_col_words = extension.dump_matrix_b_fragment( @@ -59,9 +54,7 @@ def main() -> int: True, ) matrix_b_reference_words, matrix_b_compact_words = ( - extension.compare_compact_matrix_b_col_fragment( - matrix_b_col_storage - ) + extension.compare_compact_matrix_b_col_fragment(matrix_b_col_storage) ) torch.accelerator.synchronize(device) elements = words.cpu().view(torch.float16).to(torch.int32).reshape(32, 16) @@ -69,15 +62,10 @@ def main() -> int: counts = Counter(value for lane in mapping for value in lane) multiplicities = Counter(counts.values()) matrix_b_elements = ( - matrix_b_col_words.cpu() - .view(torch.float16) - .to(torch.int32) - .reshape(32, 16) + matrix_b_col_words.cpu().view(torch.float16).to(torch.int32).reshape(32, 16) ) matrix_b_mapping = matrix_b_elements.tolist() - matrix_b_counts = Counter( - value for lane in matrix_b_mapping for value in lane - ) + matrix_b_counts = Counter(value for lane in matrix_b_mapping for value in lane) matrix_b_multiplicities = Counter(matrix_b_counts.values()) payload = { "device": str(device), @@ -94,10 +82,7 @@ def main() -> int: torch.equal(reference_words, swizzled_words) ), "swizzled_fragment_max_word_xor": int( - torch.bitwise_xor(reference_words, swizzled_words) - .abs() - .max() - .item() + torch.bitwise_xor(reference_words, swizzled_words).abs().max().item() ), "matrix_b": { "lane_elements": matrix_b_mapping, @@ -106,8 +91,7 @@ def main() -> int: str(key): value for key, value in sorted(matrix_b_multiplicities.items()) }, - "all_elements_covered": sorted(matrix_b_counts) - == list(range(256)), + "all_elements_covered": sorted(matrix_b_counts) == list(range(256)), "row_col_fragment_words_equal": bool( torch.equal(matrix_b_row_words, matrix_b_col_words) ), diff --git a/benchmarks/compare_sm70_dflash2_quality_matrix.py b/benchmarks/compare_sm70_dflash2_quality_matrix.py new file mode 100644 index 0000000000..68af86055c --- /dev/null +++ b/benchmarks/compare_sm70_dflash2_quality_matrix.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare paired multi-seed rows from the bounded DFlash2 quality matrix.""" + +from __future__ import annotations + +import argparse +import itertools +import json +import math +import random +from collections import Counter, defaultdict +from pathlib import Path +from typing import Any + +import regex as re + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--scores", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--bootstrap-samples", type=int, default=20_000) + parser.add_argument("--bootstrap-seed", type=int, default=20260826) + return parser.parse_args() + + +def _arm_name(path: Path) -> str: + match = re.search(r"-(t0|d[0-9]+)$", path.stem) + if match is None: + raise ValueError(f"cannot infer matrix arm from {path}") + return match.group(1) + + +def _ordered_arm_names(arms: dict[str, Any]) -> list[str]: + """Order the target control before progressively accelerated DFlash arms.""" + + def key(name: str) -> tuple[int, int]: + if name == "t0": + return (0, 0) + return (1, int(name[1:])) + + return sorted(arms, key=key) + + +def _percentile(values: list[float], quantile: float) -> float | None: + if not values: + return None + ordered = sorted(values) + position = (len(ordered) - 1) * quantile + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] + (ordered[upper] - ordered[lower]) * weight + + +def _paired_bootstrap_delta( + pairs: list[tuple[float, float]], + *, + samples: int, + seed: int, +) -> dict[str, float | int | None]: + if not pairs: + return {"count": 0, "mean_delta": None, "ci95": None} + deltas = [right - left for left, right in pairs] + rng = random.Random(seed) + means = [] + for _ in range(samples): + means.append(sum(rng.choice(deltas) for _ in deltas) / len(deltas)) + return { + "count": len(deltas), + "mean_delta": sum(deltas) / len(deltas), + "ci95": [_percentile(means, 0.025), _percentile(means, 0.975)], + } + + +def _mcnemar_exact(left: list[bool], right: list[bool]) -> dict[str, float | int]: + left_only = sum(a and not b for a, b in zip(left, right, strict=True)) + right_only = sum(b and not a for a, b in zip(left, right, strict=True)) + discordant = left_only + right_only + if discordant == 0: + p_value = 1.0 + else: + tail = sum( + math.comb(discordant, index) + for index in range(min(left_only, right_only) + 1) + ) / (2**discordant) + p_value = min(1.0, 2.0 * tail) + return { + "left_only": left_only, + "right_only": right_only, + "discordant": discordant, + "p_value": p_value, + } + + +def _summarize(records: list[dict[str, Any]]) -> dict[str, Any]: + passed = sum(record["passed"] for record in records) + acceptance = [ + float(record["acceptance_length"]) + for record in records + if record["acceptance_length"] is not None + ] + steady_decode = [ + float(record["steady_decode_tps"]) + for record in records + if record["steady_decode_tps"] is not None + ] + return { + "num_cases": len(records), + "passed": passed, + "score": passed / len(records) if records else None, + "finish_reasons": dict(Counter(record["finish_reason"] for record in records)), + "failure_reasons": dict( + Counter(record["reason"] for record in records if not record["passed"]) + ), + "mean_acceptance_length": ( + sum(acceptance) / len(acceptance) if acceptance else None + ), + "mean_steady_decode_tps": ( + sum(steady_decode) / len(steady_decode) if steady_decode else None + ), + } + + +def _load_arms(scores_path: Path) -> dict[str, dict[str, Any]]: + scores = json.loads(scores_path.read_text(encoding="utf-8")) + arms = {} + for scored_run in scores["runs"]: + raw_path = Path(scored_run["result"]) + raw = json.loads(raw_path.read_text(encoding="utf-8")) + if len(raw["cases"]) != len(scored_run["cases"]): + raise ValueError(f"raw/scored case count differs for {raw_path}") + + records = [] + for raw_case, scored_case in zip( + raw["cases"], scored_run["cases"], strict=True + ): + if raw_case["dataset_index"] != scored_case["dataset_index"]: + raise ValueError(f"raw/scored case order differs for {raw_path}") + request_metrics = raw_case.get("request_metrics") or {} + spec_metrics = raw_case.get("spec_decode_metrics") or {} + records.append( + { + "key": [raw_case.get("request_seed"), raw_case["dataset_index"]], + "request_seed": raw_case.get("request_seed"), + "dataset_index": raw_case["dataset_index"], + "suite": scored_case["suite"], + "suite_index": scored_case["suite_index"], + "passed": bool(scored_case["passed"]), + "reason": scored_case["reason"], + "finish_reason": raw_case["finish_reason"], + "token_hash": raw_case["token_hash"], + "output_tokens": raw_case["output_tokens"], + "acceptance_length": spec_metrics.get("acceptance_length"), + "steady_decode_tps": request_metrics.get("steady_decode_tps"), + } + ) + arm = _arm_name(raw_path) + if arm in arms: + raise ValueError(f"duplicate matrix arm {arm}") + arms[arm] = { + "result": str(raw_path), + "contract": raw["contract"], + "runtime": raw["runtime"], + "records": records, + } + return arms + + +def _record_map(arm: dict[str, Any]) -> dict[tuple[int | None, int], dict[str, Any]]: + result = {} + for record in arm["records"]: + key = tuple(record["key"]) + if key in result: + raise ValueError(f"duplicate paired key {key}") + result[key] = record + return result + + +def _compare( + left_name: str, + left_arm: dict[str, Any], + right_name: str, + right_arm: dict[str, Any], + *, + bootstrap_samples: int, + bootstrap_seed: int, +) -> dict[str, Any]: + left = _record_map(left_arm) + right = _record_map(right_arm) + if set(left) != set(right): + raise ValueError(f"paired keys differ for {left_name}/{right_name}") + keys = sorted(left) + left_pass = [left[key]["passed"] for key in keys] + right_pass = [right[key]["passed"] for key in keys] + score_pairs = [(float(a), float(b)) for a, b in zip(left_pass, right_pass)] + acceptance_pairs = [ + (float(left[key]["acceptance_length"]), float(right[key]["acceptance_length"])) + for key in keys + if left[key]["acceptance_length"] is not None + and right[key]["acceptance_length"] is not None + ] + return { + "left": left_name, + "right": right_name, + "num_cases": len(keys), + "token_hash_equal_cases": sum( + left[key]["token_hash"] == right[key]["token_hash"] for key in keys + ), + "score": _paired_bootstrap_delta( + score_pairs, + samples=bootstrap_samples, + seed=bootstrap_seed, + ), + "acceptance_length": _paired_bootstrap_delta( + acceptance_pairs, + samples=bootstrap_samples, + seed=bootstrap_seed + 1, + ), + "mcnemar_exact": _mcnemar_exact(left_pass, right_pass), + } + + +def main() -> int: + args = _parse_args() + if args.bootstrap_samples <= 0: + raise ValueError("--bootstrap-samples must be positive") + arms = _load_arms(args.scores) + + summaries = {} + for arm_name, arm in sorted(arms.items()): + by_suite: dict[str, list[dict[str, Any]]] = defaultdict(list) + by_seed: dict[str, list[dict[str, Any]]] = defaultdict(list) + for record in arm["records"]: + by_suite[record["suite"]].append(record) + by_seed[str(record["request_seed"])].append(record) + summaries[arm_name] = { + "result": arm["result"], + "aggregate": _summarize(arm["records"]), + "by_suite": { + name: _summarize(records) for name, records in sorted(by_suite.items()) + }, + "by_seed": { + name: _summarize(records) for name, records in sorted(by_seed.items()) + }, + } + + comparisons = [ + _compare( + left, + arms[left], + right, + arms[right], + bootstrap_samples=args.bootstrap_samples, + bootstrap_seed=args.bootstrap_seed, + ) + for left, right in itertools.combinations(_ordered_arm_names(arms), 2) + ] + + payload = { + "format": "sm70_dflash2_quality_matrix_v1", + "scores": str(args.scores), + "bootstrap": { + "samples": args.bootstrap_samples, + "seed": args.bootstrap_seed, + }, + "arms": summaries, + "comparisons": comparisons, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + print(json.dumps(payload, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/compare_sm70_tensor_dumps.py b/benchmarks/compare_sm70_tensor_dumps.py index 9e3de08e8a..649bf9b7c8 100644 --- a/benchmarks/compare_sm70_tensor_dumps.py +++ b/benchmarks/compare_sm70_tensor_dumps.py @@ -79,11 +79,13 @@ def _load_records(root: Path) -> dict[tuple[Any, ...], list[dict[str, Any]]]: source = payload.get("source", payload.get("layer_type", "unknown")) shape = tuple(payload.get("shape", tuple(tensor.shape))) key = (step, layer_idx, source, label, shape) - records.setdefault(key, []).append({ - "path": str(path), - "pid": payload.get("pid"), - "tensor": tensor, - }) + records.setdefault(key, []).append( + { + "path": str(path), + "pid": payload.get("pid"), + "tensor": tensor, + } + ) return records @@ -134,33 +136,37 @@ def compare_dirs(left: Path, right: Path, top: int) -> dict[str, Any]: right_items = right_records.get(key, []) step, layer_idx, source, label, shape = key if not left_items or not right_items: - missing.append({ - "step": step, - "layer_idx": layer_idx, - "source": source, - "label": label, - "shape": shape, - "left_count": len(left_items), - "right_count": len(right_items), - }) + missing.append( + { + "step": step, + "layer_idx": layer_idx, + "source": source, + "label": label, + "shape": shape, + "left_count": len(left_items), + "right_count": len(right_items), + } + ) continue _, pairs = _best_pairing(left_items, right_items) for left_idx, right_idx, stats in pairs: left_tensor = left_items[left_idx]["tensor"] right_tensor = right_items[right_idx]["tensor"] - summaries.append({ - "step": step, - "layer_idx": layer_idx, - "source": source, - "label": label, - "shape": shape, - "left_pid": left_items[left_idx]["pid"], - "right_pid": right_items[right_idx]["pid"], - "left_path": left_items[left_idx]["path"], - "right_path": right_items[right_idx]["path"], - **stats, - "rows": _row_stats(left_tensor, right_tensor), - }) + summaries.append( + { + "step": step, + "layer_idx": layer_idx, + "source": source, + "label": label, + "shape": shape, + "left_pid": left_items[left_idx]["pid"], + "right_pid": right_items[right_idx]["pid"], + "left_path": left_items[left_idx]["path"], + "right_path": right_items[right_idx]["path"], + **stats, + "rows": _row_stats(left_tensor, right_tensor), + } + ) by_worst = sorted( summaries, @@ -262,10 +268,12 @@ def _numeric_gate( if missing and not allow_missing: for item in missing[:20]: - violations.append({ - "reason": "missing_tensor_dump", - **item, - }) + violations.append( + { + "reason": "missing_tensor_dump", + **item, + } + ) for item in compared: max_abs = float(item.get("max_abs", math.inf)) @@ -280,21 +288,25 @@ def _numeric_gate( ) if item.get("shape_mismatch"): - violations.append({ - "reason": "shape_mismatch", - **_record_ref(item), - "left_shape": item.get("left_shape"), - "right_shape": item.get("right_shape"), - }) + violations.append( + { + "reason": "shape_mismatch", + **_record_ref(item), + "left_shape": item.get("left_shape"), + "right_shape": item.get("right_shape"), + } + ) continue if not math.isfinite(max_abs) or not math.isfinite(mean_abs): - violations.append({ - "reason": "non_finite_diff", - **_record_ref(item), - "max_abs": item.get("max_abs"), - "mean_abs": item.get("mean_abs"), - }) + violations.append( + { + "reason": "non_finite_diff", + **_record_ref(item), + "max_abs": item.get("max_abs"), + "mean_abs": item.get("mean_abs"), + } + ) continue item_max_abs_bound, max_abs_rule = _bound_for_item( @@ -304,21 +316,25 @@ def _numeric_gate( ) if item_max_abs_bound is None: if max_abs != 0.0: - pending.append({ - "reason": "nonzero_max_abs_without_bound", - **_record_ref(item), - "max_abs": max_abs, - }) + pending.append( + { + "reason": "nonzero_max_abs_without_bound", + **_record_ref(item), + "max_abs": max_abs, + } + ) else: checked_max_abs += 1 if max_abs > item_max_abs_bound: - violations.append({ - "reason": "max_abs_exceeds_bound", - **_record_ref(item), - "max_abs": max_abs, - "bound": item_max_abs_bound, - "bound_rule": max_abs_rule, - }) + violations.append( + { + "reason": "max_abs_exceeds_bound", + **_record_ref(item), + "max_abs": max_abs, + "bound": item_max_abs_bound, + "bound_rule": max_abs_rule, + } + ) item_mean_abs_bound, mean_abs_rule = _bound_for_item( item, @@ -328,23 +344,27 @@ def _numeric_gate( if item_mean_abs_bound is not None: checked_mean_abs += 1 if mean_abs > item_mean_abs_bound: - violations.append({ - "reason": "mean_abs_exceeds_bound", - **_record_ref(item), - "mean_abs": mean_abs, - "bound": item_mean_abs_bound, - "bound_rule": mean_abs_rule, - }) + violations.append( + { + "reason": "mean_abs_exceeds_bound", + **_record_ref(item), + "mean_abs": mean_abs, + "bound": item_mean_abs_bound, + "bound_rule": mean_abs_rule, + } + ) if min_cosine is not None and cosine is not None: checked_cosine += 1 if cosine < min_cosine: - violations.append({ - "reason": "cosine_below_bound", - **_record_ref(item), - "cosine": cosine, - "bound": min_cosine, - }) + violations.append( + { + "reason": "cosine_below_bound", + **_record_ref(item), + "cosine": cosine, + "bound": min_cosine, + } + ) if checked_max_abs == 0 and (max_abs_bound is not None or max_abs_rules): pending.append("no tensors matched any max_abs bound") diff --git a/benchmarks/kernels/benchmark_dflash2_lookup.py b/benchmarks/kernels/benchmark_dflash2_lookup.py new file mode 100644 index 0000000000..65c8300ad0 --- /dev/null +++ b/benchmarks/kernels/benchmark_dflash2_lookup.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Microbenchmark DFlash2 lookup drafting against the production UVA history.""" + +import argparse +import json +from collections.abc import Callable + +import torch + +from vllm.v1.worker.gpu.buffer_utils import UvaBuffer +from vllm.v1.worker.gpu.spec_decode.dflash2.lookup import fuse_draft, suffix_lookup + + +def _time_ms(fn: Callable[[], None], repeats: int, warmup: int) -> float: + for _ in range(warmup): + fn() + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + fn() + end.record() + end.synchronize() + return start.elapsed_time(end) / repeats + + +def _make_history(batch_size: int, context_len: int) -> UvaBuffer: + history = UvaBuffer((batch_size, context_len + 32), torch.int32) + history.cpu.random_(1000, 200000) + suffix = torch.arange(700001, 700013, dtype=torch.int32) + continuation = torch.arange(800001, 800016, dtype=torch.int32) + match_start = max(0, context_len // 2 - suffix.numel()) + for row in range(batch_size): + history.cpu[row, match_start : match_start + suffix.numel()].copy_(suffix) + history.cpu[ + row, + match_start + suffix.numel() : match_start + + suffix.numel() + + continuation.numel(), + ].copy_(continuation) + history.cpu[row, context_len - suffix.numel() : context_len].copy_(suffix) + return history + + +def _benchmark_case( + batch_size: int, + context_len: int, + repeats: int, + warmup: int, +) -> dict[str, float | int]: + device = torch.device("cuda") + k, draft_block = 15, 7 + history = _make_history(batch_size, context_len) + total_len = torch.full((batch_size,), context_len, dtype=torch.int32, device=device) + idx_mapping = torch.arange( + batch_size, dtype=torch.int32, device=device + ).repeat_interleave( + draft_block, + ) + eligible = torch.ones(batch_size, dtype=torch.int32, device=device) + lookup_tokens = torch.zeros((batch_size, k), dtype=torch.int32, device=device) + match_len = torch.zeros(batch_size, dtype=torch.int32, device=device) + valid = torch.zeros_like(match_len) + drafted = torch.zeros((batch_size, k), dtype=torch.int64, device=device) + use = torch.zeros((batch_size, k), dtype=torch.int32, device=device) + hits = torch.zeros((), dtype=torch.int64, device=device) + take_flags = torch.zeros(batch_size, dtype=torch.int32, device=device) + + def lookup_only() -> None: + suffix_lookup( + history.uva, + total_len, + idx_mapping, + eligible, + batch_size, + k, + idx_mapping_stride=draft_block, + nmin=6, + nmax=12, + out_tokens=lookup_tokens, + out_len=match_len, + out_valid=valid, + ) + + def lookup_and_fuse() -> None: + lookup_only() + fuse_draft( + drafted, + lookup_tokens, + match_len, + valid, + use, + idx_mapping, + hits, + batch_size, + k, + draft_block=draft_block, + idx_mapping_stride=draft_block, + nmin=6, + nstrong=6, + agree_min=0, + nmin_tail=4, + long_min=6, + take_flags=take_flags, + ) + + eager_ms = _time_ms(lookup_and_fuse, repeats, warmup) + graph = torch.cuda.CUDAGraph() + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + for _ in range(3): + lookup_and_fuse() + torch.cuda.current_stream().wait_stream(capture_stream) + torch.accelerator.synchronize() + with torch.cuda.graph(graph): + lookup_and_fuse() + graph_ms = _time_ms(graph.replay, repeats, warmup) + + assert torch.all(match_len == 12) + assert torch.all(valid == k) + assert torch.all(take_flags == 1) + return { + "batch_size": batch_size, + "context_len": context_len, + "uva_lookup_fuse_eager_ms": eager_ms, + "uva_lookup_fuse_graph_ms": graph_ms, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--contexts", default="1024,32768,65536,131072") + parser.add_argument("--batches", default="1,2,4,8") + parser.add_argument("--repeats", type=int, default=100) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--output", default="") + args = parser.parse_args() + + results = [ + _benchmark_case(batch, context, args.repeats, args.warmup) + for context in (int(value) for value in args.contexts.split(",")) + for batch in (int(value) for value in args.batches.split(",")) + ] + payload = { + "device": torch.cuda.get_device_name(), + "results": results, + } + rendered = json.dumps(payload, indent=2) + print(rendered) + if args.output: + with open(args.output, "w") as output_file: + output_file.write(rendered + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_dspark_markov.py b/benchmarks/kernels/benchmark_dspark_markov.py index df2185b35d..23e3202761 100644 --- a/benchmarks/kernels/benchmark_dspark_markov.py +++ b/benchmarks/kernels/benchmark_dspark_markov.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Microbenchmark DSpark's sequential Markov sampling chain.""" diff --git a/benchmarks/kernels/benchmark_dspark_metadata.py b/benchmarks/kernels/benchmark_dspark_metadata.py index 850231e984..40604d9a08 100644 --- a/benchmarks/kernels/benchmark_dspark_metadata.py +++ b/benchmarks/kernels/benchmark_dspark_metadata.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Validate DSpark query expansion and non-causal SWA metadata.""" diff --git a/benchmarks/kernels/benchmark_dspark_sm70_attention.py b/benchmarks/kernels/benchmark_dspark_sm70_attention.py index 70db0a7892..a9764bddc7 100644 --- a/benchmarks/kernels/benchmark_dspark_sm70_attention.py +++ b/benchmarks/kernels/benchmark_dspark_sm70_attention.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Replay a real DSpark packed-FP8 attention dump at N=7.""" diff --git a/benchmarks/kernels/benchmark_flash_v100_dflash2_swa.py b/benchmarks/kernels/benchmark_flash_v100_dflash2_swa.py new file mode 100644 index 0000000000..ddea51045e --- /dev/null +++ b/benchmarks/kernels/benchmark_flash_v100_dflash2_swa.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Benchmark the DFlash2 draft sliding-window paged-attention shape.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path + +import torch +from flash_attn_v100 import flash_attn_prefill_paged + + +def _time_ms(fn, *, warmup: int, reps: int) -> list[float]: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + samples = [] + for _ in range(reps): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + return samples + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--seq-lens", + type=int, + nargs="+", + default=[1024, 32768, 131072, 261888], + ) + parser.add_argument("--q-len", type=int, default=8) + parser.add_argument("--num-query-heads", type=int, default=8) + parser.add_argument("--num-kv-heads", type=int, default=2) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--block-size", type=int, default=16) + parser.add_argument("--window-left", type=int, default=2048) + parser.add_argument("--window-right", type=int, default=2048) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--reps", type=int, default=100) + parser.add_argument("--output", type=Path) + parser.add_argument("--tensor-output", type=Path) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if torch.cuda.get_device_capability() != (7, 0): + raise RuntimeError("This benchmark targets SM70") + + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.float16 + max_seq_len = max(args.seq_lens) + num_blocks = (max_seq_len + args.block_size - 1) // args.block_size + query = torch.randn( + 1, + args.q_len, + args.num_query_heads, + args.head_dim, + device=device, + dtype=dtype, + ) + key_cache = torch.randn( + num_blocks, + args.block_size, + args.num_kv_heads, + args.head_dim, + device=device, + dtype=dtype, + ) + value_cache = torch.randn_like(key_cache) + full_block_table = torch.arange( + num_blocks, device=device, dtype=torch.int32 + ).unsqueeze(0) + output = torch.empty_like(query) + tensor_outputs: dict[int, torch.Tensor] = {} + results = [] + + for seq_len in args.seq_lens: + active_blocks = (seq_len + args.block_size - 1) // args.block_size + block_table = full_block_table[:, :active_blocks] + seq_lens = torch.tensor([seq_len], device=device, dtype=torch.int32) + + def run( + block_table: torch.Tensor = block_table, + seq_lens: torch.Tensor = seq_lens, + ) -> None: + flash_attn_prefill_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + out=output, + causal=False, + window_size=(args.window_left, args.window_right), + ) + + samples = _time_ms(run, warmup=args.warmup, reps=args.reps) + run() + torch.cuda.synchronize() + tensor_outputs[seq_len] = output.detach().cpu().clone() + results.append( + { + "seq_len": seq_len, + "median_ms": statistics.median(samples), + "mean_ms": statistics.mean(samples), + "min_ms": min(samples), + "max_ms": max(samples), + "five_draft_layers_median_ms": 5 * statistics.median(samples), + } + ) + + report = { + "device": torch.cuda.get_device_name(), + "torch": torch.__version__, + "shape": { + "q_len": args.q_len, + "num_query_heads": args.num_query_heads, + "num_kv_heads": args.num_kv_heads, + "head_dim": args.head_dim, + "block_size": args.block_size, + "causal": False, + "window_size": [args.window_left, args.window_right], + }, + "warmup": args.warmup, + "reps": args.reps, + "results": results, + } + rendered = json.dumps(report, indent=2) + print(rendered) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + if args.tensor_output is not None: + args.tensor_output.parent.mkdir(parents=True, exist_ok=True) + torch.save(tensor_outputs, args.tensor_output) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_dflash2_grouped_verify.py b/benchmarks/kernels/benchmark_sm70_dflash2_grouped_verify.py new file mode 100644 index 0000000000..1545c512e0 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_dflash2_grouped_verify.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Microbenchmark the SM70 q8/q16 grouped DFlash2 verifier.""" + +from __future__ import annotations + +import argparse +import json +import math +from collections.abc import Callable +from functools import partial + +import torch + + +def _make_case( + query_len: int, prefix_len: int, page_size: int +) -> tuple[torch.Tensor, ...]: + total_len = prefix_len + query_len + logical_pages = math.ceil(total_len / page_size) + physical_pages = logical_pages + 3 + source = torch.randn( + (physical_pages, 2, page_size, 1, 256), + dtype=torch.float16, + device="cuda", + ).mul_(0.25) + cache = source.to(torch.float8_e5m2).view(torch.uint8) + key_cache, value_cache = cache.unbind(1) + block_table = torch.randperm(physical_pages, dtype=torch.int32, device="cuda")[ + :logical_pages + ].view(1, -1) + query = torch.randn((query_len, 6, 256), dtype=torch.float16, device="cuda").mul_( + 0.25 + ) + seq_lens = torch.tensor([total_len], dtype=torch.int32, device="cuda") + return query, key_cache, value_cache, block_table, seq_lens + + +def _time_cuda(fn: Callable[[], object], warmup: int, repeats: int) -> float: + for _ in range(warmup): + fn() + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + fn() + end.record() + end.synchronize() + return start.elapsed_time(end) / repeats + + +def _capture(fn: Callable[[], object]) -> torch.cuda.CUDAGraph: + fn() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + return graph + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--query-lens", type=int, nargs="+", default=[8, 16]) + parser.add_argument( + "--prefix-lens", type=int, nargs="+", default=[1024, 32768, 128000] + ) + parser.add_argument("--page-size", type=int, default=3296) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--repeats", type=int, default=100) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if torch.cuda.get_device_capability() != (7, 0): + raise RuntimeError("This benchmark requires an SM70 GPU") + + import flash_attn_v100 + + torch.manual_seed(20260827) + results = [] + for query_len in args.query_lens: + for prefix_len in args.prefix_lens: + query, key_cache, value_cache, block_table, seq_lens = _make_case( + query_len, prefix_len, args.page_size + ) + grouped_out = torch.empty_like(query) + independent_out = torch.empty_like(query) + decode_block_table = block_table.repeat(query_len, 1).contiguous() + decode_seq_lens = torch.arange( + prefix_len + 1, + prefix_len + query_len + 1, + dtype=torch.int32, + device=query.device, + ) + + grouped = partial( + flash_attn_v100.flash_attn_grouped_verify_paged, + query, + key_cache, + value_cache, + block_table, + seq_lens, + out=grouped_out, + one_pass=True, + ) + independent = partial( + flash_attn_v100.flash_attn_decode_paged_xqa, + query, + key_cache, + value_cache, + decode_block_table, + decode_seq_lens, + out=independent_out, + kv_cache_dtype="fp8_e5m2", + max_seq_len_hint=prefix_len + query_len, + workspace_seq_capacity_hint=prefix_len + query_len, + ) + + grouped() + independent() + torch.accelerator.synchronize() + difference = grouped_out.float().sub(independent_out.float()).abs() + grouped_eager_ms = _time_cuda(grouped, args.warmup, args.repeats) + independent_eager_ms = _time_cuda(independent, args.warmup, args.repeats) + grouped_graph = _capture(grouped) + independent_graph = _capture(independent) + grouped_graph_ms = _time_cuda( + grouped_graph.replay, args.warmup, args.repeats + ) + independent_graph_ms = _time_cuda( + independent_graph.replay, args.warmup, args.repeats + ) + results.append( + { + "query_len": query_len, + "prefix_len": prefix_len, + "page_size": args.page_size, + "grouped_eager_ms": grouped_eager_ms, + "independent_eager_ms": independent_eager_ms, + "eager_speedup": independent_eager_ms / grouped_eager_ms, + "grouped_graph_ms": grouped_graph_ms, + "independent_graph_ms": independent_graph_ms, + "graph_speedup": independent_graph_ms / grouped_graph_ms, + "max_abs_vs_independent": difference.max().item(), + "mean_abs_vs_independent": difference.mean().item(), + } + ) + + print( + json.dumps( + { + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "results": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_fp8_graph_capture.py b/benchmarks/kernels/benchmark_sm70_fp8_graph_capture.py index ba73f21797..d31121dd36 100644 --- a/benchmarks/kernels/benchmark_sm70_fp8_graph_capture.py +++ b/benchmarks/kernels/benchmark_sm70_fp8_graph_capture.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Validate an uncached SM70 FP8 GEMM shape under CUDA Graph capture.""" diff --git a/benchmarks/verify_sm70_artifact.py b/benchmarks/verify_sm70_artifact.py index dfca12798f..81dde92497 100644 --- a/benchmarks/verify_sm70_artifact.py +++ b/benchmarks/verify_sm70_artifact.py @@ -75,18 +75,14 @@ ("sm70_turbomind_policy.awq_turbomind_effective", True), ) -AWQ_TURBOMIND_DENSE_REQUIRED_LOGS = ( - "SM70 AWQ TurboMind dense path enabled", -) +AWQ_TURBOMIND_DENSE_REQUIRED_LOGS = ("SM70 AWQ TurboMind dense path enabled",) FP8_TURBOMIND_DENSE_JSON_EXPECTATIONS = ( ("sm70_turbomind_policy.fp8_turbomind_effective", True), ("sm70_turbomind_policy.fp8_dequant_fallback_effective", True), ) -FP8_TURBOMIND_DENSE_REQUIRED_LOGS = ( - "SM70 FP8 TurboMind W8A16 dense path enabled", -) +FP8_TURBOMIND_DENSE_REQUIRED_LOGS = ("SM70 FP8 TurboMind W8A16 dense path enabled",) FP8_0DOT3_DENSE_DEQUANT_JSON_EXPECTATIONS = ( ("sm70_turbomind_policy.fp8_turbomind_effective", False), @@ -105,8 +101,7 @@ ("sm70_turbomind_policy.accepted_awq_moe_default_policy", True), ("sm70_turbomind_policy.awq_moe_disable_effective", False), ("sm70_turbomind_policy.awq_moe_batched_gemm_effective", False), - ("sm70_turbomind_policy.awq_moe_legacy_single_token_compact_effective", - False), + ("sm70_turbomind_policy.awq_moe_legacy_single_token_compact_effective", False), ("sm70_moe_policy.single_token_unpermute_fastpath_effective", True), ) @@ -121,8 +116,7 @@ ("sm70_turbomind_policy.accepted_awq_moe_default_policy", True), ("sm70_turbomind_policy.awq_moe_disable_effective", False), ("sm70_turbomind_policy.awq_moe_batched_gemm_effective", True), - ("sm70_turbomind_policy.awq_moe_legacy_single_token_compact_effective", - False), + ("sm70_turbomind_policy.awq_moe_legacy_single_token_compact_effective", False), ) AWQ_MOE_BATCHED_SAFE_ROUTE_REQUIRED_LOGS = ( @@ -138,8 +132,7 @@ ("sm70_turbomind_policy.awq_moe_0dot3_baseline_policy", True), ("sm70_turbomind_policy.awq_moe_disable_effective", False), ("sm70_turbomind_policy.awq_moe_batched_gemm_effective", True), - ("sm70_turbomind_policy.awq_moe_legacy_single_token_compact_effective", - True), + ("sm70_turbomind_policy.awq_moe_legacy_single_token_compact_effective", True), ) AWQ_MOE_0DOT3_BASELINE_ROUTE_REQUIRED_LOGS = ( @@ -153,9 +146,7 @@ ("sm70_turbomind_policy.fp8_moe_batched_gemm_effective", True), ) -FP8_MOE_SAFE_ROUTE_REQUIRED_LOGS = ( - "SM70 FP8 MoE TurboMind batched path enabled", -) +FP8_MOE_SAFE_ROUTE_REQUIRED_LOGS = ("SM70 FP8 MoE TurboMind batched path enabled",) FP8_MOE_0DOT3_FALLBACK_ROUTE_JSON_EXPECTATIONS = ( ("sm70_turbomind_policy.fp8_moe_0dot3_dequant_fallback_policy", True), @@ -281,9 +272,7 @@ def _get_json_path(payload: Any, path: str) -> Any: def _parse_expectation(raw: str) -> tuple[str, Any]: if "=" not in raw: - raise ValueError( - f"Expected PATH=VALUE for --expect-json, got {raw!r}" - ) + raise ValueError(f"Expected PATH=VALUE for --expect-json, got {raw!r}") path, expected = raw.split("=", 1) if not path: raise ValueError(f"Expected non-empty PATH in {raw!r}") @@ -339,17 +328,13 @@ def _awq_moe_safe_tune_gate_for_meta(meta: Any) -> dict[str, Any]: } raw = tune_policy.get("VLLM_SM70_AWQ_TUNE_SMALL_SHAPES") - tune_safe_default = tune_policy.get( - "awq_moe_safe_default_selector_effective" - ) + tune_safe_default = tune_policy.get("awq_moe_safe_default_selector_effective") turbomind_safe_default = turbomind_policy.get( "awq_moe_safe_default_selector_effective" ) pinned = raw == "0" source_safe_unset = ( - raw is None - and tune_safe_default is True - and turbomind_safe_default is True + raw is None and tune_safe_default is True and turbomind_safe_default is True ) return { "passed": pinned or source_safe_unset, @@ -409,10 +394,12 @@ def _awq_moe_safe_tune_gate(payload: Any) -> dict[str, Any]: return { "passed": False, "location": None, - "checks": [{ - "passed": False, - "reason": "no top-level or compare metadata policies found", - }], + "checks": [ + { + "passed": False, + "reason": "no top-level or compare metadata policies found", + } + ], } @@ -529,8 +516,7 @@ def _model_quality_failures(payload: Any) -> list[str]: failures = [] if gate.get("label") != "model-pass": failures.append( - "model_quality_gate.label is not model-pass " - f"({gate.get('label')!r})" + f"model_quality_gate.label is not model-pass ({gate.get('label')!r})" ) if gate.get("default_acceptance") != "model-level gate passed": failures.append( @@ -586,8 +572,7 @@ def _inventory_complete_failures(payload: Any) -> list[str]: failures = [] if gate.get("label") != "inventory-complete": failures.append( - "inventory_gate.label is not inventory-complete " - f"({gate.get('label')!r})" + f"inventory_gate.label is not inventory-complete ({gate.get('label')!r})" ) if gate.get("complete") is not True: failures.append("inventory_gate.complete is not true") @@ -595,8 +580,7 @@ def _inventory_complete_failures(payload: Any) -> list[str]: old_only_env_vars = gate.get("old_only_env_vars") if old_only_env_vars: failures.append( - "inventory_gate.old_only_env_vars is not empty: " - f"{old_only_env_vars!r}" + f"inventory_gate.old_only_env_vars is not empty: {old_only_env_vars!r}" ) if gate.get("old_only_env_vars_count") != 0: failures.append( @@ -647,11 +631,13 @@ def _verify( present = location is not None if not present: missing_policies.append(policy) - policy_results.append({ - "policy": policy, - "present": present, - "location": location, - }) + policy_results.append( + { + "policy": policy, + "present": present, + "location": location, + } + ) json_results = [] failed_json_expectations = [] @@ -678,10 +664,12 @@ def _verify( present = text in logs if not present: missing_log_texts.append(text) - log_results.append({ - "text": text, - "present": present, - }) + log_results.append( + { + "text": text, + "present": present, + } + ) log_regex_results = [] missing_log_regexes = [] @@ -689,10 +677,12 @@ def _verify( present = re.search(pattern, logs) is not None if not present: missing_log_regexes.append(pattern) - log_regex_results.append({ - "pattern": pattern, - "present": present, - }) + log_regex_results.append( + { + "pattern": pattern, + "present": present, + } + ) rejected_log_results = [] rejected_log_texts = [] @@ -700,10 +690,12 @@ def _verify( present = text in logs if present: rejected_log_texts.append(text) - rejected_log_results.append({ - "text": text, - "present": present, - }) + rejected_log_results.append( + { + "text": text, + "present": present, + } + ) rejected_log_regex_results = [] rejected_log_regexes = [] @@ -711,10 +703,12 @@ def _verify( present = re.search(pattern, logs) is not None if present: rejected_log_regexes.append(pattern) - rejected_log_regex_results.append({ - "pattern": pattern, - "present": present, - }) + rejected_log_regex_results.append( + { + "pattern": pattern, + "present": present, + } + ) decode_timing_failures = ( _decode_timing_failures(payload) if require_decode_timing else [] @@ -723,8 +717,7 @@ def _verify( _model_quality_failures(payload) if require_model_quality_pass else [] ) inventory_complete_failures = ( - _inventory_complete_failures(payload) - if require_inventory_complete else [] + _inventory_complete_failures(payload) if require_inventory_complete else [] ) passed = not ( @@ -888,9 +881,7 @@ def _build_parser() -> argparse.ArgumentParser: parser.add_argument( "--require-awq-turbomind-dense", action="store_true", - help=( - "Require SM70 AWQ TurboMind dense policy and its runtime route log." - ), + help=("Require SM70 AWQ TurboMind dense policy and its runtime route log."), ) parser.add_argument( "--require-fp8-turbomind-dense", @@ -1017,9 +1008,7 @@ def main() -> int: try: payload = _read_json(args.json) - expectations = [ - _parse_expectation(raw) for raw in args.expect_json - ] + expectations = [_parse_expectation(raw) for raw in args.expect_json] log_expectations = list(args.expect_log) log_regex_expectations = list(args.expect_log_regex) rejected_log_expectations = list(args.reject_log) @@ -1034,9 +1023,7 @@ def main() -> int: log_expectations.extend(FULL_FLASH_V100_REQUIRED_LOGS) log_expectations.extend(SM70_FP8_KV_CACHE_ROUTE_REQUIRED_LOGS) rejected_log_expectations.extend(FULL_FLASH_V100_REJECTED_LOGS) - rejected_log_expectations.extend( - SM70_FP8_KV_CACHE_ROUTE_REJECTED_LOGS - ) + rejected_log_expectations.extend(SM70_FP8_KV_CACHE_ROUTE_REJECTED_LOGS) if args.require_turbomind_default_policy: expectations.extend(TURBOMIND_DEFAULT_JSON_EXPECTATIONS) if args.require_awq_turbomind_dense: @@ -1057,27 +1044,17 @@ def main() -> int: if args.require_awq_moe_batched_safe_route: expectations.extend(AWQ_MOE_BATCHED_SAFE_ROUTE_JSON_EXPECTATIONS) log_expectations.extend(AWQ_MOE_BATCHED_SAFE_ROUTE_REQUIRED_LOGS) - rejected_log_expectations.extend( - AWQ_MOE_BATCHED_SAFE_ROUTE_REJECTED_LOGS - ) + rejected_log_expectations.extend(AWQ_MOE_BATCHED_SAFE_ROUTE_REJECTED_LOGS) if args.require_awq_moe_0dot3_baseline_route: expectations.extend(AWQ_MOE_0DOT3_BASELINE_ROUTE_JSON_EXPECTATIONS) - log_expectations.extend( - AWQ_MOE_0DOT3_BASELINE_ROUTE_REQUIRED_LOGS - ) + log_expectations.extend(AWQ_MOE_0DOT3_BASELINE_ROUTE_REQUIRED_LOGS) if args.require_fp8_moe_safe_route: expectations.extend(FP8_MOE_SAFE_ROUTE_JSON_EXPECTATIONS) log_expectations.extend(FP8_MOE_SAFE_ROUTE_REQUIRED_LOGS) if args.require_fp8_moe_0dot3_fallback_route: - expectations.extend( - FP8_MOE_0DOT3_FALLBACK_ROUTE_JSON_EXPECTATIONS - ) - log_expectations.extend( - FP8_MOE_0DOT3_FALLBACK_ROUTE_REQUIRED_LOGS - ) - rejected_log_expectations.extend( - FP8_MOE_0DOT3_FALLBACK_ROUTE_REJECTED_LOGS - ) + expectations.extend(FP8_MOE_0DOT3_FALLBACK_ROUTE_JSON_EXPECTATIONS) + log_expectations.extend(FP8_MOE_0DOT3_FALLBACK_ROUTE_REQUIRED_LOGS) + rejected_log_expectations.extend(FP8_MOE_0DOT3_FALLBACK_ROUTE_REJECTED_LOGS) extra_required_policies = list(args.require_policy) if args.require_sm70_breakable_graph_route: extra_required_policies.append("sm70_graph_policy") @@ -1085,12 +1062,8 @@ def main() -> int: log_expectations.extend(SM70_BREAKABLE_GRAPH_ROUTE_REQUIRED_LOGS) if args.require_sm70_flash_v100_0dot3_compile_graph: extra_required_policies.append("sm70_graph_policy") - expectations.extend( - SM70_FLASH_V100_0DOT3_COMPILE_GRAPH_JSON_EXPECTATIONS - ) - log_expectations.extend( - SM70_FLASH_V100_0DOT3_COMPILE_GRAPH_REQUIRED_LOGS - ) + expectations.extend(SM70_FLASH_V100_0DOT3_COMPILE_GRAPH_JSON_EXPECTATIONS) + log_expectations.extend(SM70_FLASH_V100_0DOT3_COMPILE_GRAPH_REQUIRED_LOGS) rejected_log_expectations.extend( SM70_FLASH_V100_0DOT3_COMPILE_GRAPH_REJECTED_LOGS ) @@ -1098,19 +1071,17 @@ def main() -> int: extra_required_policies.append("sm70_comm_policy") expectations.extend(SM70_CUSTOM_ALLREDUCE_ROUTE_JSON_EXPECTATIONS) log_expectations.extend(SM70_CUSTOM_ALLREDUCE_ROUTE_REQUIRED_LOGS) - rejected_log_expectations.extend( - SM70_CUSTOM_ALLREDUCE_ROUTE_REJECTED_LOGS - ) + rejected_log_expectations.extend(SM70_CUSTOM_ALLREDUCE_ROUTE_REJECTED_LOGS) if args.require_sm70_all_reduce_sum2_route: extra_required_policies.append("sm70_comm_policy") expectations.extend(SM70_ALL_REDUCE_SUM2_ROUTE_JSON_EXPECTATIONS) log_expectations.extend(SM70_ALL_REDUCE_SUM2_ROUTE_REQUIRED_LOGS) - log_regex_expectations.extend( - SM70_ALL_REDUCE_SUM2_ROUTE_REQUIRED_REGEXES - ) - required_policies = [] if args.require_inventory_complete else list( - dict.fromkeys( - list(DEFAULT_REQUIRED_POLICIES) + extra_required_policies + log_regex_expectations.extend(SM70_ALL_REDUCE_SUM2_ROUTE_REQUIRED_REGEXES) + required_policies = ( + [] + if args.require_inventory_complete + else list( + dict.fromkeys(list(DEFAULT_REQUIRED_POLICIES) + extra_required_policies) ) ) result = _verify( @@ -1127,10 +1098,7 @@ def main() -> int: require_model_quality_pass=args.require_model_quality_pass, require_inventory_complete=args.require_inventory_complete, ) - if ( - args.require_awq_moe_safe_route - or args.require_awq_moe_batched_safe_route - ): + if args.require_awq_moe_safe_route or args.require_awq_moe_batched_safe_route: awq_moe_safe_tune_gate = _awq_moe_safe_tune_gate(payload) result["awq_moe_safe_tune_gate"] = awq_moe_safe_tune_gate if not awq_moe_safe_tune_gate["passed"]: diff --git a/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch b/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch index d6d65439a5..8809667dc1 100644 --- a/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch +++ b/cmake/patches/sm70_flash_attn_d256_gqa_arch.patch @@ -1211,8 +1211,8 @@ index 0000000..197453f + constexpr int kHeadsQ = 6; + constexpr int kHeadsKV = 1; + constexpr int kTail = 8000; -+ constexpr int kMinTotalKV = 40000; -+ constexpr int kMaxTotalKV = 128000; ++ constexpr int kMinTotalKV = 16000; ++ constexpr int kMaxTotalKV = 256000; + constexpr int kTotalKVStep = 8000; + constexpr int kBlockN = 8192; + @@ -1232,7 +1232,7 @@ index 0000000..197453f + const int total_kv = static_cast(k.size(1)); + TORCH_CHECK(total_kv >= kMinTotalKV && total_kv <= kMaxTotalKV + && total_kv % kTotalKVStep == 0, -+ "SM70 GQA architecture requires KV in [40000, 128000] " ++ "SM70 GQA architecture requires KV in [16000, 256000] " + "with an 8000-token step, got ", + total_kv); + const int prefix = total_kv - kTail; diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index f98aef77f5..a19ffe9e94 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -70,6 +70,7 @@ constexpr int kSm70Tp4PushAllreduceSentinelByte = 0x7f; constexpr size_t kSm70Tp4PushAllreduceBytes = 8 * kSm70GemmaRmsNormHiddenSize * sizeof(half); constexpr size_t kSm70Tp4PushAllreduce8KiBBytes = 4096 * sizeof(half); +constexpr size_t kSm70Tp4PushAllreduceQwen4ExpBytes = 2560 * sizeof(half); constexpr size_t kSm70Tp4PushAllreduceSignalBytes = ((kSm70Tp4PushAllreduceBlocks * sizeof(uint32_t) + 127) / 128) * 128; constexpr size_t kSm70Tp4PushAllreduceBufferBytes = @@ -81,7 +82,10 @@ inline int sm70_tp4_push_allreduce_blocks(size_t bytes) { if (bytes == kSm70Tp4PushAllreduceBytes) { return kSm70Tp4PushAllreduceBlocks; } - return bytes == kSm70Tp4PushAllreduce8KiBBytes ? 4 : 0; + if (bytes == kSm70Tp4PushAllreduce8KiBBytes) { + return 4; + } + return bytes == kSm70Tp4PushAllreduceQwen4ExpBytes ? 3 : 0; } inline int sm70_gemma_rms_norm_threads() { diff --git a/csrc/moe/marlin_moe_wna16/sm70_marlin_fp8_gemm.cu b/csrc/moe/marlin_moe_wna16/sm70_marlin_fp8_gemm.cu index 1329091679..290e531181 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_fp8_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_fp8_gemm.cu @@ -35,22 +35,25 @@ class Sm70MoeFp8IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; - static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, - "SM70 Marlin MoE FP8 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, - "SM70 Marlin MoE FP8 IteratorB expects one contiguous iteration per " - "64-column quant tile."); + using Fragment = cutlass::Array; + static_assert( + Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, + "SM70 Marlin MoE FP8 IteratorB expects CTA_N in {64, 128, 256}."); + static_assert( + ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, + "SM70 Marlin MoE FP8 IteratorB expects one contiguous iteration per " + "64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, "SM70 Marlin MoE FP8 IteratorB expects 64-column deltas."); - static_assert(ThreadMap::kElementsPerAccess == kFp8ValuesPerAccess, - "SM70 Marlin MoE FP8 IteratorB expects two packed FP8 words per " - "access."); - static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin MoE U8-family IteratorB expects one or more strided " - "iterations."); + static_assert( + ThreadMap::kElementsPerAccess == kFp8ValuesPerAccess, + "SM70 Marlin MoE FP8 IteratorB expects two packed FP8 words per " + "access."); + static_assert( + ThreadMap::Iterations::kStrided >= 1, + "SM70 Marlin MoE U8-family IteratorB expects one or more strided " + "iterations."); static constexpr int kQweightWordStrideWords = kPackedMacroN / kQuantTileN; struct Params { @@ -83,12 +86,11 @@ class Sm70MoeFp8IteratorB { public: CUTLASS_DEVICE Sm70MoeFp8IteratorB(Params const& params, uint32_t const* qweight, - half const* scales, half const*, int thread_id, - int expert, int k_offset, int n_offset) + half const* scales, half const*, int thread_id, + int expert, int k_offset, int n_offset) : qweight_(qweight), - scales_expert_base_(scales + - (expert >= 0 ? expert : 0) * - params.num_groups * params.size_n), + scales_expert_base_(scales + (expert >= 0 ? expert : 0) * + params.num_groups * params.size_n), params_(params), thread_offset_(ThreadMap::initial_offset(thread_id)), expert_(expert), @@ -130,9 +132,7 @@ class Sm70MoeFp8IteratorB { } CUTLASS_DEVICE - int metadata_group_offset(int group) const { - return group * params_.size_n; - } + int metadata_group_offset(int group) const { return group * params_.size_n; } CUTLASS_DEVICE int scale_group(int logical_k) const { @@ -155,8 +155,8 @@ class Sm70MoeFp8IteratorB { CUTLASS_DEVICE void cache_metadata_lane_vectors(int c, int group, int cache_n) const { int const metadata_offset = metadata_group_offset(group) + cache_n; - half2 const* scale_vec = reinterpret_cast( - scales_expert_base_ + metadata_offset); + half2 const* scale_vec = + reinterpret_cast(scales_expert_base_ + metadata_offset); half2* scale_cache = cached_scales_ + c * 4; scale_cache[0] = scale_vec[0]; scale_cache[1] = scale_vec[1]; @@ -167,8 +167,8 @@ class Sm70MoeFp8IteratorB { CUTLASS_DEVICE void cache_metadata_vector_words(int c, int group, int cache_n) const { int const metadata_offset = metadata_group_offset(group) + cache_n; - uint4 const scale_words = *reinterpret_cast( - scales_expert_base_ + metadata_offset); + uint4 const scale_words = + *reinterpret_cast(scales_expert_base_ + metadata_offset); half2 const* scale_vec = reinterpret_cast(&scale_words); half2* scale_cache = cached_scales_ + c * 4; scale_cache[0] = scale_vec[0]; @@ -181,9 +181,8 @@ class Sm70MoeFp8IteratorB { void cache_current_group_metadata(int group) const { CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; if constexpr (kUseMetadataVectorWords) { cache_metadata_vector_words(c, group, cache_n); } else { @@ -244,8 +243,9 @@ class Sm70MoeFp8IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE FP8 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE FP8 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_offset_); uint32_t const qword1 = load_qword_vector<1>( @@ -267,9 +267,8 @@ class Sm70MoeFp8IteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; if constexpr (kGroupSize != -1) { cache_current_group_metadata(scale_group(logical_k_s)); } @@ -277,10 +276,9 @@ class Sm70MoeFp8IteratorB { expert_qweight_base_offset() + qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - uint4 const qwords0 = - load_qword_vector<4>(qweight_ + qweight_base_s); - uint4 const qwords1 = load_qword_vector<4>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint4 const qwords0 = load_qword_vector<4>(qweight_ + qweight_base_s); + uint4 const qwords1 = load_qword_vector<4>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -294,22 +292,17 @@ class Sm70MoeFp8IteratorB { half2* frag_vec = reinterpret_cast(frag.data() + frag_base); marlin::dequant( static_cast(qword0), deq); - frag_vec[0] = - __hmul2(deq[0], scale_vec[0]); - frag_vec[1] = - __hmul2(deq[1], scale_vec[1]); + frag_vec[0] = __hmul2(deq[0], scale_vec[0]); + frag_vec[1] = __hmul2(deq[1], scale_vec[1]); marlin::dequant( static_cast(qword1), deq); - frag_vec[2] = - __hmul2(deq[0], scale_vec[2]); - frag_vec[3] = - __hmul2(deq[1], scale_vec[3]); + frag_vec[2] = __hmul2(deq[0], scale_vec[2]); + frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - uint2 const qwords0 = - load_qword_vector<2>(qweight_ + qweight_base_s); - uint2 const qwords1 = load_qword_vector<2>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint2 const qwords0 = load_qword_vector<2>(qweight_ + qweight_base_s); + uint2 const qwords1 = load_qword_vector<2>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -323,20 +316,17 @@ class Sm70MoeFp8IteratorB { half2* frag_vec = reinterpret_cast(frag.data() + frag_base); marlin::dequant( static_cast(qword0), deq); - frag_vec[0] = - __hmul2(deq[0], scale_vec[0]); - frag_vec[1] = - __hmul2(deq[1], scale_vec[1]); + frag_vec[0] = __hmul2(deq[0], scale_vec[0]); + frag_vec[1] = __hmul2(deq[1], scale_vec[1]); marlin::dequant( static_cast(qword1), deq); - frag_vec[2] = - __hmul2(deq[0], scale_vec[2]); - frag_vec[3] = - __hmul2(deq[1], scale_vec[3]); + frag_vec[2] = __hmul2(deq[0], scale_vec[2]); + frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE FP8 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE FP8 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_s); uint32_t const qword1 = load_qword_vector<1>( @@ -366,8 +356,7 @@ class Sm70MoeFp8IteratorB { return; } - if constexpr (kGroupSize != -1 && - ThreadMap::Iterations::kStrided == 1) { + if constexpr (kGroupSize != -1 && ThreadMap::Iterations::kStrided == 1) { int const first_logical_k = k_offset_ + thread_offset_.strided(); cache_current_group_metadata(scale_group(first_logical_k)); } @@ -386,18 +375,17 @@ struct Sm70MoeFp8GemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = - Sm70MoeFp8IteratorB; + using IteratorB = Sm70MoeFp8IteratorB; }; -template using Sm70MoeFp8GemmTraits = - Sm70MarlinMoeGemmTraits< - Sm70MoeFp8GemmSpec, CtaM, CtaN, CtaK, - Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; + Sm70MarlinMoeGemmTraits, CtaM, + CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>; template struct Sm70MoeFp8Launcher { @@ -419,16 +407,16 @@ struct Sm70MoeFp8Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - using Traits = Sm70MoeFp8GemmTraits; + using Traits = + Sm70MoeFp8GemmTraits; return launch_sm70_marlin_moe_gemm( a, c, b_q_weight, b_scales, b_zeros, global_scale, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, - top_k, mul_topk_weights, size_m, size_n, size_k, requested_split_k); + expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, + mul_topk_weights, size_m, size_n, size_k, requested_split_k); } }; @@ -446,9 +434,10 @@ torch::Tensor sm70_marlin_fp8_gemm( auto const params = sm70_marlin_moe_auto_stage_params( "fp8_e4m3", group_size, moe_block_size, top_k, size_m, size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE fp8_e4m3", geometry); - validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE fp8_e4m3", geometry, - size_n); + validate_sm70_marlin_moe_stage_cta_geometry_supported( + "SM70 Marlin MoE fp8_e4m3", geometry); + validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE fp8_e4m3", + geometry, size_n); TORCH_CHECK(size_k % geometry.cta_k == 0, "SM70 Marlin MoE fp8_e4m3 requires K divisible by CTA_K=", geometry.cta_k, ". Got K=", size_k, "."); @@ -457,17 +446,45 @@ torch::Tensor sm70_marlin_fp8_gemm( auto empty_float = torch::empty( {0}, torch::TensorOptions().dtype(at::kFloat).device(a.device())); if (params.use_metadata_vector_words) { - Sm70MoeFp8Launcher const launcher{ - a, c, b_q_weight, b_scales, empty_half, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; - return dispatch_sm70_marlin_moe_fp8_geometry(launcher, geometry, params.packed_macro_n, group_size); + Sm70MoeFp8Launcher const launcher{a, + c, + b_q_weight, + b_scales, + empty_half, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; + return dispatch_sm70_marlin_moe_fp8_geometry( + launcher, geometry, params.packed_macro_n, group_size); } - Sm70MoeFp8Launcher const launcher{ - a, c, b_q_weight, b_scales, empty_half, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; - return dispatch_sm70_marlin_moe_fp8_geometry(launcher, geometry, params.packed_macro_n, group_size); + Sm70MoeFp8Launcher const launcher{a, + c, + b_q_weight, + b_scales, + empty_half, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; + return dispatch_sm70_marlin_moe_fp8_geometry( + launcher, geometry, params.packed_macro_n, group_size); } } // namespace marlin_moe_wna16 diff --git a/csrc/moe/marlin_moe_wna16/sm70_marlin_gemm.cuh b/csrc/moe/marlin_moe_wna16/sm70_marlin_gemm.cuh index fdc516a3a6..2cd0b71f90 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_gemm.cuh +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_gemm.cuh @@ -27,22 +27,22 @@ namespace marlin_moe_wna16 { -using marlin::sm70::Sm70CtaGeometry; -using marlin::sm70::Sm70MarlinAutoParams; -using marlin::sm70::Sm70MarlinMoeAutoParamsContext; -using marlin::sm70::Sm70SplitKPartition; -using marlin::sm70::Sm70MarlinMmaPipelined; -using marlin::sm70::Sm70WarpShape; using marlin::sm70::configure_sm70_dynamic_smem; using marlin::sm70::kQuantTileK; using marlin::sm70::kQuantTileN; +using marlin::sm70::sm70_active_split_k; using marlin::sm70::sm70_marlin_any_auto_env_is_set; using marlin::sm70::sm70_marlin_auto_packed_macro_n; using marlin::sm70::sm70_marlin_auto_params_from_env; using marlin::sm70::sm70_marlin_cta_geometry_is_supported; using marlin::sm70::sm70_marlin_default_auto_params; -using marlin::sm70::sm70_active_split_k; using marlin::sm70::sm70_splitk_partition; +using marlin::sm70::Sm70CtaGeometry; +using marlin::sm70::Sm70MarlinAutoParams; +using marlin::sm70::Sm70MarlinMmaPipelined; +using marlin::sm70::Sm70MarlinMoeAutoParamsContext; +using marlin::sm70::Sm70SplitKPartition; +using marlin::sm70::Sm70WarpShape; using marlin::sm70::validate_sm70_marlin_auto_params; inline constexpr char const* kSupportedSm70MarlinMoeCtaGeometries = @@ -55,9 +55,9 @@ inline constexpr char const* kSupportedSm70MarlinMoeCtaGeometries = "geometries are dense-only."; inline bool sm70_marlin_moe_auto_env_is_set() { - return sm70_marlin_any_auto_env_is_set( - "SM70_MARLIN_MOE_CTA_GEOMETRY", "SM70_MARLIN_MOE_SPLIT_K", - "SM70_MARLIN_MOE_METADATA_CACHE"); + return sm70_marlin_any_auto_env_is_set("SM70_MARLIN_MOE_CTA_GEOMETRY", + "SM70_MARLIN_MOE_SPLIT_K", + "SM70_MARLIN_MOE_METADATA_CACHE"); } inline Sm70MarlinAutoParams sm70_marlin_moe_auto_params_from_env( @@ -68,16 +68,14 @@ inline Sm70MarlinAutoParams sm70_marlin_moe_auto_params_from_env( } inline bool sm70_marlin_moe_try_select_quanttrio_qwen3_6_35b_a3b_awq_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4") != 0 || - ctx.group_size != 128 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4") != 0 || ctx.group_size != 128 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -152,8 +150,7 @@ inline bool sm70_marlin_moe_try_select_quanttrio_qwen3_6_35b_a3b_awq_params( } return set_params({64, 128, 32, 8, 32, 32, 32}, 1, false); } - if (ctx.moe_block_size == 64 && - (ctx.size_n == 256 || ctx.size_n == 1024)) { + if (ctx.moe_block_size == 64 && (ctx.size_n == 256 || ctx.size_n == 1024)) { return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } @@ -162,16 +159,14 @@ inline bool sm70_marlin_moe_try_select_quanttrio_qwen3_6_35b_a3b_awq_params( } inline bool sm70_marlin_moe_try_select_quanttrio_qwen3_5_122b_a10b_awq_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4") != 0 || - ctx.group_size != 128 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4") != 0 || ctx.group_size != 128 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -278,16 +273,14 @@ inline bool sm70_marlin_moe_try_select_quanttrio_qwen3_5_122b_a10b_awq_params( } inline bool sm70_marlin_moe_try_select_quanttrio_minimax_m2_7_awq_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4") != 0 || - ctx.group_size != 128 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4") != 0 || ctx.group_size != 128 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -378,17 +371,16 @@ inline bool sm70_marlin_moe_try_select_quanttrio_minimax_m2_7_awq_params( return false; } -inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool +sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_params( + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint8b128") != 0 || - ctx.group_size != 32 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint8b128") != 0 || ctx.group_size != 32 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -400,57 +392,57 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_param // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 128) { - if (ctx.size_m <= 8) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - if (ctx.size_m <= 256) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + if (ctx.size_m <= 256) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 256) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 1024) { - if (ctx.size_m <= 8) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 256) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,128,8,32,64,32}, 1, true); - } - return set_params({32,128,64,8,32,32,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 128, 64, 8, 32, 32, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 512) { - if (ctx.size_m <= 1) { - return set_params({32,256,32,4,32,64,32}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,256,64,8,32,64,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 2048) { - if (ctx.size_m <= 1) { - return set_params({32,256,32,4,32,64,32}, 2, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 2, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } } @@ -460,13 +452,13 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_param // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 1024) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 2048) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } } @@ -476,13 +468,13 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_param // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 1024) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 2048) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } } @@ -492,33 +484,33 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_param // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 128) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 256) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 1024) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 256) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 512) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 2048) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } } @@ -527,16 +519,14 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_5_122b_a10b_awq_8bit_param } inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_4bit_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4") != 0 || - ctx.group_size != 32 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4") != 0 || ctx.group_size != 32 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -548,63 +538,63 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_4bit_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - if (ctx.size_m <= 8) { - return set_params({32,64,32,4,32,32,16}, 1, false); - } - if (ctx.size_m <= 256) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, false); + } + if (ctx.size_m <= 256) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - if (ctx.size_m <= 8) { - return set_params({32,64,32,4,32,32,16}, 1, true); - } - if (ctx.size_m <= 256) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, true); + } + if (ctx.size_m <= 256) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - if (ctx.size_m <= 8) { - return set_params({32,128,32,4,32,64,16}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 128, 32, 4, 32, 64, 16}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, false); - } - if (ctx.size_m <= 32) { - return set_params({32,128,64,8,32,64,16}, 2, false); - } - return set_params({32,128,128,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, false); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 64, 8, 32, 64, 16}, 2, false); + } + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,128,8,32,64,32}, 1, true); - } - return set_params({32,256,64,8,32,64,32}, 1, false); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, false); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - if (ctx.size_m <= 1) { - return set_params({32,256,32,4,32,64,32}, 4, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 4, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } } @@ -614,16 +604,16 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_4bit_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - if (ctx.size_m <= 32) { - return set_params({32,256,64,8,32,64,32}, 1, false); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 32) { + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, false); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } } @@ -633,13 +623,13 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_4bit_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({32,256,64,8,32,64,32}, 1, true); + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } } @@ -649,36 +639,36 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_4bit_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - if (ctx.size_m <= 2048) { - return set_params({64,128,32,4,64,64,16}, 1, false); - } - return set_params({64,128,32,8,32,32,32}, 1, true); + if (ctx.size_m <= 2048) { + return set_params({64, 128, 32, 4, 64, 64, 16}, 1, false); + } + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } } @@ -686,17 +676,16 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_4bit_params( return false; } -inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_nvfp4_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool +sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_nvfp4_params( + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "nvfp4") != 0 || - ctx.group_size != 16 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "nvfp4") != 0 || ctx.group_size != 16 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -708,63 +697,63 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_nvfp4_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - if (ctx.size_m <= 8) { - return set_params({32,128,64,4,32,64,32}, 1, true); - } - if (ctx.size_m <= 256) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 128, 64, 4, 32, 64, 32}, 1, true); + } + if (ctx.size_m <= 256) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - if (ctx.size_m <= 8) { - return set_params({32,128,32,4,32,32,32}, 1, false); - } - if (ctx.size_m <= 256) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, false); + if (ctx.size_m <= 8) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, false); + } + if (ctx.size_m <= 256) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - if (ctx.size_m <= 8) { - return set_params({32,128,64,4,32,64,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 128, 64, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,64,8,32,32,32}, 4, true); - } - return set_params({32,128,128,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 64, 32, 4, 32, 32, 16}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 64, 8, 32, 32, 32}, 4, true); + } + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,128,8,32,64,32}, 1, true); - } - return set_params({32,256,64,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 4, false); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 4, false); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } } @@ -774,16 +763,16 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_nvfp4_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({32,128,32,4,32,32,32}, 1, false); + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, false); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - if (ctx.size_m <= 32) { - return set_params({32,256,64,8,32,64,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 32) { + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } } @@ -793,13 +782,13 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_nvfp4_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({32,128,32,4,32,32,32}, 1, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({32,256,64,8,32,64,32}, 1, false); + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, false); } } } @@ -809,45 +798,45 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_nvfp4_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - if (ctx.size_m <= 16384) { - return set_params({64,64,32,4,32,32,32}, 1, true); - } - return set_params({64,256,32,4,64,64,32}, 1, false); + if (ctx.size_m <= 16384) { + return set_params({64, 64, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - if (ctx.size_m <= 16384) { - return set_params({64,256,32,4,64,64,32}, 1, false); - } - return set_params({64,256,32,4,64,64,32}, 1, true); + if (ctx.size_m <= 16384) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - if (ctx.size_m <= 2048) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({64,128,32,8,32,32,32}, 1, true); + if (ctx.size_m <= 2048) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - if (ctx.size_m <= 2048) { - return set_params({64,256,32,4,64,64,32}, 1, true); - } - return set_params({64,256,64,8,64,64,32}, 1, false); + if (ctx.size_m <= 2048) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); + } + return set_params({64, 256, 64, 8, 64, 64, 32}, 1, false); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } } @@ -855,17 +844,16 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_6_35b_a3b_awq_nvfp4_params return false; } -inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_4bit_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool +sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_4bit_params( + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4b8") != 0 || - ctx.group_size != 32 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4b8") != 0 || ctx.group_size != 32 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -877,63 +865,63 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_4bit_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - if (ctx.size_m <= 10) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - if (ctx.size_m <= 320) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 10) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + if (ctx.size_m <= 320) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - if (ctx.size_m <= 10) { - return set_params({32,128,64,4,32,64,32}, 1, true); - } - if (ctx.size_m <= 320) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 10) { + return set_params({32, 128, 64, 4, 32, 64, 32}, 1, true); + } + if (ctx.size_m <= 320) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,64,8,32,64,16}, 2, true); - } - return set_params({32,128,128,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 64, 32, 4, 32, 32, 16}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 64, 8, 32, 64, 16}, 2, true); + } + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,128,8,32,64,32}, 1, true); - } - return set_params({32,256,64,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - if (ctx.size_m <= 1) { - return set_params({32,256,32,4,32,64,32}, 4, true); - } - if (ctx.size_m <= 32) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 4, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } } @@ -943,13 +931,13 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_4bit_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } } @@ -959,23 +947,23 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_4bit_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - return set_params({64,128,32,8,32,32,32}, 1, false); + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, false); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } } @@ -985,33 +973,33 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_4bit_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - return set_params({64,128,32,4,32,64,32}, 1, true); + return set_params({64, 128, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } } @@ -1019,17 +1007,16 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_4bit_params return false; } -inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_8bit_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool +sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_8bit_params( + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint8b128") != 0 || - ctx.group_size != 32 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint8b128") != 0 || ctx.group_size != 32 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -1041,63 +1028,63 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_8bit_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - if (ctx.size_m <= 10) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - if (ctx.size_m <= 320) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 10) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + if (ctx.size_m <= 320) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - if (ctx.size_m <= 10) { - return set_params({32,256,64,8,32,64,32}, 1, true); - } - if (ctx.size_m <= 320) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 10) { + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); + } + if (ctx.size_m <= 320) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,64,8,32,64,16}, 2, true); - } - return set_params({32,128,128,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 64, 32, 4, 32, 32, 16}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 64, 8, 32, 64, 16}, 2, true); + } + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - if (ctx.size_m <= 1) { - return set_params({32,256,32,4,32,64,32}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,128,8,32,64,32}, 1, true); - } - return set_params({32,256,64,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - if (ctx.size_m <= 1) { - return set_params({32,128,128,8,32,64,32}, 1, true); - } - if (ctx.size_m <= 32) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } } @@ -1107,13 +1094,13 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_8bit_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({32,256,32,4,32,64,32}, 1, true); + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } } @@ -1123,23 +1110,23 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_8bit_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - return set_params({64,128,32,8,32,32,32}, 1, false); + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, false); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } } @@ -1149,33 +1136,33 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_8bit_params // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - return set_params({64,128,32,8,32,32,32}, 1, false); + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, false); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } } @@ -1184,16 +1171,14 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_qwen3_coder_next_awq_8bit_params } inline bool sm70_marlin_moe_try_select_nvidia_glm_4_7_nvfp4_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "nvfp4") != 0 || - ctx.group_size != 16 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "nvfp4") != 0 || ctx.group_size != 16 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -1205,63 +1190,63 @@ inline bool sm70_marlin_moe_try_select_nvidia_glm_4_7_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 5120) { if (ctx.size_k == 192) { - if (ctx.size_m <= 8) { - return set_params({32,128,32,4,32,32,32}, 1, false); - } - if (ctx.size_m <= 256) { - return set_params({32,128,32,4,32,32,32}, 1, false); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, false); + } + if (ctx.size_m <= 256) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, false); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 5120) { if (ctx.size_k == 384) { - if (ctx.size_m <= 8) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - if (ctx.size_m <= 256) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + if (ctx.size_m <= 256) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 5120) { if (ctx.size_k == 1536) { - if (ctx.size_m <= 8) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, false); + if (ctx.size_m <= 8) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, false); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 5120) { if (ctx.size_n == 384) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,32,4,32,64,16}, 2, true); - } - return set_params({32,128,32,4,32,64,16}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 32, 4, 32, 64, 16}, 2, true); + } + return set_params({32, 128, 32, 4, 32, 64, 16}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 5120) { if (ctx.size_n == 768) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,64,16}, 4, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,32,4,32,64,16}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 4, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 64, 16}, 4, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 32, 4, 32, 64, 16}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 4, true); } } if (ctx.top_k == 8 && ctx.size_k == 5120) { if (ctx.size_n == 3072) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); - } - return set_params({32,128,32,4,32,32,32}, 2, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 2, true); } } } @@ -1271,19 +1256,19 @@ inline bool sm70_marlin_moe_try_select_nvidia_glm_4_7_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 5120) { if (ctx.size_k == 1536) { - if (ctx.size_m <= 256) { - return set_params({32,256,64,8,32,64,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 256) { + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 5120) { if (ctx.size_n == 3072) { - if (ctx.size_m <= 32) { - return set_params({32,128,32,4,32,32,32}, 4, true); - } - return set_params({32,256,64,8,32,64,32}, 1, true); + if (ctx.size_m <= 32) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 4, true); + } + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } } @@ -1293,13 +1278,13 @@ inline bool sm70_marlin_moe_try_select_nvidia_glm_4_7_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 5120) { if (ctx.size_k == 1536) { - return set_params({32,256,64,8,32,64,32}, 1, true); + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 5120) { if (ctx.size_n == 3072) { - return set_params({32,256,64,8,32,64,32}, 1, true); + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } } @@ -1309,42 +1294,42 @@ inline bool sm70_marlin_moe_try_select_nvidia_glm_4_7_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 5120) { if (ctx.size_k == 192) { - if (ctx.size_m <= 16384) { - return set_params({64,256,32,4,64,64,32}, 1, true); - } - return set_params({64,256,32,4,64,64,32}, 1, false); + if (ctx.size_m <= 16384) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 5120) { if (ctx.size_k == 384) { - if (ctx.size_m <= 16384) { - return set_params({64,256,32,4,64,64,32}, 1, false); - } - return set_params({64,256,32,4,64,64,32}, 1, true); + if (ctx.size_m <= 16384) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 5120) { if (ctx.size_k == 1536) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 5120) { if (ctx.size_n == 384) { - if (ctx.size_m <= 2048) { - return set_params({64,128,32,8,32,32,32}, 1, false); - } - return set_params({64,128,32,8,32,32,32}, 1, true); + if (ctx.size_m <= 2048) { + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, false); + } + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 5120) { if (ctx.size_n == 768) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 5120) { if (ctx.size_n == 3072) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } } @@ -1353,16 +1338,14 @@ inline bool sm70_marlin_moe_try_select_nvidia_glm_4_7_nvfp4_params( } inline bool sm70_marlin_moe_try_select_nvidia_minimax_m2_7_nvfp4_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "nvfp4") != 0 || - ctx.group_size != 16 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "nvfp4") != 0 || ctx.group_size != 16 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -1374,63 +1357,63 @@ inline bool sm70_marlin_moe_try_select_nvidia_minimax_m2_7_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 192) { - if (ctx.size_m <= 8) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - if (ctx.size_m <= 256) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + if (ctx.size_m <= 256) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 384) { - if (ctx.size_m <= 8) { - return set_params({32,256,32,4,32,64,32}, 1, true); - } - if (ctx.size_m <= 256) { - return set_params({32,256,32,4,32,64,32}, 1, false); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); + } + if (ctx.size_m <= 256) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, false); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 1536) { - if (ctx.size_m <= 8) { - return set_params({32,128,32,4,32,64,16}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 128, 32, 4, 32, 64, 16}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 384) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, false); - } - if (ctx.size_m <= 32) { - return set_params({32,128,32,4,32,64,16}, 2, true); - } - return set_params({32,128,32,4,32,64,16}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, false); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 32, 4, 32, 64, 16}, 2, true); + } + return set_params({32, 128, 32, 4, 32, 64, 16}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 768) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 4, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,32,4,32,64,16}, 1, true); - } - return set_params({32,128,32,4,32,64,16}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 4, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 32, 4, 32, 64, 16}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 64, 16}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 3072) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 4, true); - } - return set_params({32,256,64,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 4, true); + } + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } } @@ -1440,19 +1423,19 @@ inline bool sm70_marlin_moe_try_select_nvidia_minimax_m2_7_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 1536) { - if (ctx.size_m <= 256) { - return set_params({32,256,64,8,32,64,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 256) { + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 3072) { - if (ctx.size_m <= 32) { - return set_params({32,256,64,8,32,64,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 32) { + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } } @@ -1462,13 +1445,13 @@ inline bool sm70_marlin_moe_try_select_nvidia_minimax_m2_7_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 1536) { - return set_params({32,256,64,8,32,64,32}, 1, true); + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 3072) { - return set_params({32,256,64,8,32,64,32}, 1, false); + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, false); } } } @@ -1478,45 +1461,45 @@ inline bool sm70_marlin_moe_try_select_nvidia_minimax_m2_7_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 192) { - if (ctx.size_m <= 16384) { - return set_params({64,256,32,4,64,64,32}, 1, true); - } - return set_params({64,256,32,4,64,64,32}, 1, false); + if (ctx.size_m <= 16384) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 384) { - if (ctx.size_m <= 16384) { - return set_params({64,256,32,4,64,64,32}, 1, false); - } - return set_params({64,256,32,4,64,64,32}, 1, true); + if (ctx.size_m <= 16384) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 3072) { if (ctx.size_k == 1536) { - if (ctx.size_m <= 16384) { - return set_params({64,256,32,4,64,64,32}, 1, true); - } - return set_params({64,256,32,4,64,64,32}, 1, false); + if (ctx.size_m <= 16384) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 384) { - if (ctx.size_m <= 2048) { - return set_params({64,128,32,8,32,32,32}, 1, true); - } - return set_params({64,128,32,8,32,32,32}, 1, false); + if (ctx.size_m <= 2048) { + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, true); + } + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, false); } } if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 768) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 3072) { if (ctx.size_n == 3072) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } } @@ -1525,16 +1508,14 @@ inline bool sm70_marlin_moe_try_select_nvidia_minimax_m2_7_nvfp4_params( } inline bool sm70_marlin_moe_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "nvfp4") != 0 || - ctx.group_size != 16 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "nvfp4") != 0 || ctx.group_size != 16 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -1546,63 +1527,63 @@ inline bool sm70_marlin_moe_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - if (ctx.size_m <= 8) { - return set_params({32,64,64,4,32,32,32}, 1, true); - } - if (ctx.size_m <= 256) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); + } + if (ctx.size_m <= 256) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - if (ctx.size_m <= 8) { - return set_params({32,64,64,4,32,32,32}, 1, true); - } - if (ctx.size_m <= 256) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); + } + if (ctx.size_m <= 256) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - if (ctx.size_m <= 8) { - return set_params({32,128,64,4,32,64,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 8) { + return set_params({32, 128, 64, 4, 32, 64, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,64,8,32,32,32}, 4, true); - } - return set_params({32,128,128,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 64, 32, 4, 32, 32, 16}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 64, 8, 32, 32, 32}, 4, true); + } + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,128,8,32,64,32}, 1, true); - } - return set_params({32,256,64,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 4, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 4, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } } @@ -1612,16 +1593,16 @@ inline bool sm70_marlin_moe_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({32,128,32,4,32,32,32}, 1, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - if (ctx.size_m <= 32) { - return set_params({32,256,64,8,32,64,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 32) { + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } } @@ -1631,13 +1612,13 @@ inline bool sm70_marlin_moe_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({32,128,32,4,32,32,32}, 1, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({32,256,64,8,32,64,32}, 1, true); + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } } @@ -1647,48 +1628,48 @@ inline bool sm70_marlin_moe_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params( // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - if (ctx.size_m <= 16384) { - return set_params({64,64,32,4,32,32,32}, 1, false); - } - return set_params({64,256,32,4,64,64,32}, 1, false); + if (ctx.size_m <= 16384) { + return set_params({64, 64, 32, 4, 32, 32, 32}, 1, false); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - if (ctx.size_m <= 16384) { - return set_params({64,256,32,4,64,64,32}, 1, false); - } - return set_params({64,256,32,4,64,64,32}, 1, true); + if (ctx.size_m <= 16384) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } // top_k=8 if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - if (ctx.size_m <= 2048) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({64,128,32,8,32,32,32}, 1, true); + if (ctx.size_m <= 2048) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, true); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - if (ctx.size_m <= 2048) { - return set_params({64,256,32,4,64,64,32}, 1, true); - } - return set_params({64,256,64,8,64,64,32}, 1, false); + if (ctx.size_m <= 2048) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); + } + return set_params({64, 256, 64, 8, 64, 64, 32}, 1, false); } } if (ctx.top_k == 8 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - if (ctx.size_m <= 2048) { - return set_params({64,256,32,4,64,64,32}, 1, false); - } - return set_params({64,256,32,4,64,64,32}, 1, true); + if (ctx.size_m <= 2048) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } } @@ -1696,17 +1677,16 @@ inline bool sm70_marlin_moe_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params( return false; } -inline bool sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool +sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_params( + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "nvfp4") != 0 || - ctx.group_size != 16 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "nvfp4") != 0 || ctx.group_size != 16 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -1718,69 +1698,69 @@ inline bool sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_ // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - if (ctx.size_m <= 10) { - return set_params({32,128,32,4,32,32,32}, 1, false); - } - if (ctx.size_m <= 320) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 10) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, false); + } + if (ctx.size_m <= 320) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - if (ctx.size_m <= 10) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - if (ctx.size_m <= 320) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,256,32,4,32,64,32}, 1, true); + if (ctx.size_m <= 10) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + if (ctx.size_m <= 320) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - if (ctx.size_m <= 10) { - return set_params({32,256,64,8,32,64,32}, 1, true); - } - if (ctx.size_m <= 320) { - return set_params({32,128,32,4,32,32,32}, 1, true); - } - return set_params({32,128,32,4,32,32,32}, 1, false); + if (ctx.size_m <= 10) { + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); + } + if (ctx.size_m <= 320) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, false); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 8, false); - } - if (ctx.size_m <= 32) { - return set_params({32,128,64,8,32,64,16}, 2, true); - } - return set_params({32,128,128,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 64, 32, 4, 32, 32, 16}, 8, false); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 64, 8, 32, 64, 16}, 2, true); + } + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); - } - if (ctx.size_m <= 32) { - return set_params({32,128,128,8,32,64,32}, 1, true); - } - return set_params({32,256,64,8,32,64,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); + } + return set_params({32, 256, 64, 8, 32, 64, 32}, 1, true); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - if (ctx.size_m <= 1) { - return set_params({32,128,128,8,32,64,32}, 1, true); - } - if (ctx.size_m <= 32) { - return set_params({32,256,32,4,32,64,32}, 1, false); - } - return set_params({32,128,32,4,32,32,32}, 1, true); + if (ctx.size_m <= 1) { + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); + } + if (ctx.size_m <= 32) { + return set_params({32, 256, 32, 4, 32, 64, 32}, 1, false); + } + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } } @@ -1790,13 +1770,13 @@ inline bool sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_ // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - return set_params({32,128,32,4,32,32,32}, 1, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({32,128,32,4,32,32,32}, 1, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } } } @@ -1806,23 +1786,23 @@ inline bool sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_ // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - return set_params({64,128,32,8,32,32,32}, 1, false); + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, false); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - return set_params({64,256,32,4,64,64,32}, 1, false); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); } } } @@ -1832,36 +1812,36 @@ inline bool sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_ // top_k=1 if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 64) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 128) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } if (ctx.top_k == 1 && ctx.size_n == 2048) { if (ctx.size_k == 512) { - if (ctx.size_m <= 20480) { - return set_params({64,256,32,4,64,64,32}, 1, false); - } - return set_params({64,256,32,4,64,64,32}, 1, true); + if (ctx.size_m <= 20480) { + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, false); + } + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } // top_k=10 if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 128) { - return set_params({64,128,32,8,32,32,32}, 1, true); + return set_params({64, 128, 32, 8, 32, 32, 32}, 1, true); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 256) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } if (ctx.top_k == 10 && ctx.size_k == 2048) { if (ctx.size_n == 1024) { - return set_params({64,256,32,4,64,64,32}, 1, true); + return set_params({64, 256, 32, 4, 64, 64, 32}, 1, true); } } } @@ -1870,16 +1850,14 @@ inline bool sm70_marlin_moe_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_ } inline bool sm70_marlin_moe_try_select_cyankiwi_minimax_m2_7_awq_4bit_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4b8") != 0 || - ctx.group_size != 32 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4b8") != 0 || ctx.group_size != 32 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -1995,16 +1973,14 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_minimax_m2_7_awq_4bit_params( } inline bool sm70_marlin_moe_try_select_cyankiwi_glm_4_7_awq_4bit_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4b8") != 0 || - ctx.group_size != 32 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4b8") != 0 || ctx.group_size != 32 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -2108,16 +2084,14 @@ inline bool sm70_marlin_moe_try_select_cyankiwi_glm_4_7_awq_4bit_params( } inline bool sm70_marlin_moe_try_select_quanttrio_glm_4_7_awq_params( - Sm70MarlinMoeAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinMoeAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4") != 0 || - ctx.group_size != 128 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4") != 0 || ctx.group_size != 128 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -2206,20 +2180,21 @@ inline bool sm70_marlin_moe_try_select_quanttrio_glm_4_7_awq_params( } inline Sm70MarlinAutoParams sm70_marlin_moe_auto_stage_params( - char const* quant_format, int64_t group_size, - int64_t moe_block_size, int64_t top_k, int64_t size_m, - int64_t size_n, int64_t size_k) { + char const* quant_format, int64_t group_size, int64_t moe_block_size, + int64_t top_k, int64_t size_m, int64_t size_n, int64_t size_k) { Sm70MarlinMoeAutoParamsContext const ctx{ - quant_format, group_size, moe_block_size, top_k, size_m, size_n, size_k, - sm70_marlin_auto_packed_macro_n(size_n)}; + quant_format, group_size, + moe_block_size, top_k, + size_m, size_n, + size_k, sm70_marlin_auto_packed_macro_n(size_n)}; if (sm70_marlin_moe_auto_env_is_set()) { return sm70_marlin_moe_auto_params_from_env(ctx); } Sm70MarlinAutoParams params{}; - if (sm70_marlin_moe_try_select_quanttrio_qwen3_6_35b_a3b_awq_params( - ctx, params)) { + if (sm70_marlin_moe_try_select_quanttrio_qwen3_6_35b_a3b_awq_params(ctx, + params)) { validate_sm70_marlin_auto_params("MoE", params); return params; } @@ -2228,8 +2203,8 @@ inline Sm70MarlinAutoParams sm70_marlin_moe_auto_stage_params( validate_sm70_marlin_auto_params("MoE", params); return params; } - if (sm70_marlin_moe_try_select_quanttrio_minimax_m2_7_awq_params( - ctx, params)) { + if (sm70_marlin_moe_try_select_quanttrio_minimax_m2_7_awq_params(ctx, + params)) { validate_sm70_marlin_auto_params("MoE", params); return params; } @@ -2258,18 +2233,17 @@ inline Sm70MarlinAutoParams sm70_marlin_moe_auto_stage_params( validate_sm70_marlin_auto_params("MoE", params); return params; } - if (sm70_marlin_moe_try_select_nvidia_glm_4_7_nvfp4_params( - ctx, params)) { + if (sm70_marlin_moe_try_select_nvidia_glm_4_7_nvfp4_params(ctx, params)) { validate_sm70_marlin_auto_params("MoE", params); return params; } - if (sm70_marlin_moe_try_select_nvidia_minimax_m2_7_nvfp4_params( - ctx, params)) { + if (sm70_marlin_moe_try_select_nvidia_minimax_m2_7_nvfp4_params(ctx, + params)) { validate_sm70_marlin_auto_params("MoE", params); return params; } - if (sm70_marlin_moe_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params( - ctx, params)) { + if (sm70_marlin_moe_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params(ctx, + params)) { validate_sm70_marlin_auto_params("MoE", params); return params; } @@ -2283,13 +2257,12 @@ inline Sm70MarlinAutoParams sm70_marlin_moe_auto_stage_params( validate_sm70_marlin_auto_params("MoE", params); return params; } - if (sm70_marlin_moe_try_select_cyankiwi_glm_4_7_awq_4bit_params( - ctx, params)) { + if (sm70_marlin_moe_try_select_cyankiwi_glm_4_7_awq_4bit_params(ctx, + params)) { validate_sm70_marlin_auto_params("MoE", params); return params; } - if (sm70_marlin_moe_try_select_quanttrio_glm_4_7_awq_params( - ctx, params)) { + if (sm70_marlin_moe_try_select_quanttrio_glm_4_7_awq_params(ctx, params)) { validate_sm70_marlin_auto_params("MoE", params); return params; } @@ -2303,46 +2276,44 @@ inline bool sm70_marlin_moe_cta_geometry_is_supported( sm70_marlin_cta_geometry_is_supported(geometry); } -inline void validate_sm70_marlin_moe_stage_cta_geometry_supported(char const* op_name, - Sm70CtaGeometry geometry) { +inline void validate_sm70_marlin_moe_stage_cta_geometry_supported( + char const* op_name, Sm70CtaGeometry geometry) { TORCH_CHECK(sm70_marlin_moe_cta_geometry_is_supported(geometry), "Unsupported SM70 Marlin MoE CTA geometry for ", op_name, ": ", - geometry.cta_m, "x", geometry.cta_n, "x", geometry.cta_k, - "x", geometry.warps, "x", geometry.warp_m, "x", - geometry.warp_n, "x", geometry.warp_k, - ". Supported geometries are ", + geometry.cta_m, "x", geometry.cta_n, "x", geometry.cta_k, "x", + geometry.warps, "x", geometry.warp_m, "x", geometry.warp_n, "x", + geometry.warp_k, ". Supported geometries are ", kSupportedSm70MarlinMoeCtaGeometries, "."); } -inline void validate_sm70_marlin_moe_stage_cta_n_alignment(char const* op_name, - Sm70CtaGeometry geometry, - int64_t size_n) { +inline void validate_sm70_marlin_moe_stage_cta_n_alignment( + char const* op_name, Sm70CtaGeometry geometry, int64_t size_n) { TORCH_CHECK( size_n % geometry.cta_n == 0 && size_n % kQuantTileN == 0, "SM70 Marlin MoE requires size_n divisible by both CTA_N and 64 for ", - op_name, " with CTA geometry ", geometry.cta_m, "x", geometry.cta_n, - "x", geometry.cta_k, "x", geometry.warps, "x", geometry.warp_m, - "x", geometry.warp_n, "x", geometry.warp_k, ". Got size_n = ", - size_n, "."); + op_name, " with CTA geometry ", geometry.cta_m, "x", geometry.cta_n, "x", + geometry.cta_k, "x", geometry.warps, "x", geometry.warp_m, "x", + geometry.warp_n, "x", geometry.warp_k, ". Got size_n = ", size_n, "."); } template -torch::Tensor dispatch_sm70_marlin_moe_group_size( - Launcher const& launcher, int64_t group_size, char const* quant_name) { +torch::Tensor dispatch_sm70_marlin_moe_group_size(Launcher const& launcher, + int64_t group_size, + char const* quant_name) { switch (group_size) { case -1: - return launcher.template operator()(); + return launcher.template operator()(); case 32: - return launcher.template operator()(); + return launcher.template operator()(); case 64: - return launcher.template operator()(); + return launcher.template operator()(); case 128: - return launcher.template operator()(); + return launcher.template operator()(); default: TORCH_CHECK(false, "SM70 Marlin MoE ", quant_name, " supports only group_size -1, 32, 64, or 128. Got ", @@ -2357,17 +2328,17 @@ torch::Tensor dispatch_sm70_marlin_moe_group_size_with_group32( Launcher const& launcher, int64_t group_size, char const* quant_name) { switch (group_size) { case -1: - return launcher.template operator()(); + return launcher.template operator()(); case 32: - return launcher.template operator()(); + return launcher.template operator()(); case 64: - return launcher.template operator()(); + return launcher.template operator()(); case 128: - return launcher.template operator()(); + return launcher.template operator()(); default: TORCH_CHECK(false, "SM70 Marlin MoE ", quant_name, " supports only group_size -1, 32, 64, or 128. Got ", @@ -2378,15 +2349,15 @@ torch::Tensor dispatch_sm70_marlin_moe_group_size_with_group32( template -torch::Tensor dispatch_sm70_marlin_moe_fp8_group_size( - Launcher const& launcher, int64_t group_size) { +torch::Tensor dispatch_sm70_marlin_moe_fp8_group_size(Launcher const& launcher, + int64_t group_size) { switch (group_size) { case -1: - return launcher.template operator()(); + return launcher.template operator()(); case 128: - return launcher.template operator()(); + return launcher.template operator()(); default: TORCH_CHECK(false, "SM70 Marlin MoE fp8_e4m3 supports only group_size -1 " @@ -2448,10 +2419,9 @@ struct Sm70MarlinMoeFp8GroupSizeDispatchLauncher { template torch::Tensor operator()() const { - return dispatch_sm70_marlin_moe_fp8_group_size( - inner, group_size); + return dispatch_sm70_marlin_moe_fp8_group_size< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, PackedMacroN>(inner, + group_size); } }; @@ -2465,59 +2435,59 @@ struct Sm70MarlinMoeFixedGroupSizeDispatchLauncher { int WarpK, int PackedMacroN> torch::Tensor operator()() const { return dispatch_sm70_marlin_moe_fixed_group_size< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN>(inner, group_size, quant_name); + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>( + inner, group_size, quant_name); } }; template -torch::Tensor dispatch_sm70_marlin_moe_cta_geometry( - Launcher const& launcher, Sm70CtaGeometry geometry, int packed_macro_n, - char const* quant_name) { -#define DISPATCH_SM70_MOE_PACKED_MACRO_N(CM, CN, CK, W, WM, WN, WK, PMN) \ +torch::Tensor dispatch_sm70_marlin_moe_cta_geometry(Launcher const& launcher, + Sm70CtaGeometry geometry, + int packed_macro_n, + char const* quant_name) { +#define DISPATCH_SM70_MOE_PACKED_MACRO_N(CM, CN, CK, W, WM, WN, WK, PMN) \ if (packed_macro_n == PMN) { \ return launcher.template operator()(); \ } -#define DISPATCH_SM70_MOE_GEOMETRY(CM, CN, CK, W, WM, WN, WK, PMN) \ - if (geometry.cta_m == CM && geometry.cta_n == CN && \ - geometry.cta_k == CK && geometry.warps == W && \ - geometry.warp_m == WM && geometry.warp_n == WN && \ - geometry.warp_k == WK) { \ - DISPATCH_SM70_MOE_PACKED_MACRO_N(CM, CN, CK, W, WM, WN, WK, PMN) \ - } - -#define FOR_EACH_SM70_MOE_GEOMETRY(M) \ - M(32, 64, 32, 4, 32, 32, 16, 128) \ - M(32, 64, 32, 4, 32, 32, 16, 256) \ - M(32, 64, 64, 4, 32, 32, 32, 128) \ - M(32, 64, 64, 4, 32, 32, 32, 256) \ - M(32, 128, 32, 4, 32, 32, 32, 128) \ - M(32, 128, 32, 4, 32, 32, 32, 256) \ - M(32, 128, 32, 4, 32, 64, 16, 128) \ - M(32, 128, 32, 4, 32, 64, 16, 256) \ - M(32, 128, 64, 4, 32, 64, 32, 128) \ - M(32, 128, 64, 4, 32, 64, 32, 256) \ - M(32, 128, 64, 8, 32, 32, 32, 128) \ - M(32, 128, 64, 8, 32, 32, 32, 256) \ - M(32, 128, 64, 8, 32, 64, 16, 128) \ - M(32, 128, 64, 8, 32, 64, 16, 256) \ - M(32, 128, 128, 8, 32, 64, 32, 128) \ - M(32, 128, 128, 8, 32, 64, 32, 256) \ - M(32, 256, 32, 4, 32, 64, 32, 256) \ - M(32, 256, 64, 8, 32, 64, 32, 256) \ - M(64, 64, 32, 4, 32, 32, 32, 64) \ - M(64, 64, 32, 4, 32, 32, 32, 128) \ - M(64, 64, 32, 4, 32, 32, 32, 256) \ - M(64, 128, 32, 4, 32, 64, 32, 128) \ - M(64, 128, 32, 4, 32, 64, 32, 256) \ - M(64, 128, 32, 4, 64, 64, 16, 128) \ - M(64, 128, 32, 4, 64, 64, 16, 256) \ - M(64, 128, 32, 8, 32, 32, 32, 128) \ - M(64, 128, 32, 8, 32, 32, 32, 256) \ - M(64, 128, 64, 4, 64, 64, 32, 128) \ - M(64, 128, 64, 4, 64, 64, 32, 256) \ - M(64, 256, 32, 4, 64, 64, 32, 256) \ +#define DISPATCH_SM70_MOE_GEOMETRY(CM, CN, CK, W, WM, WN, WK, PMN) \ + if (geometry.cta_m == CM && geometry.cta_n == CN && geometry.cta_k == CK && \ + geometry.warps == W && geometry.warp_m == WM && geometry.warp_n == WN && \ + geometry.warp_k == WK) { \ + DISPATCH_SM70_MOE_PACKED_MACRO_N(CM, CN, CK, W, WM, WN, WK, PMN) \ + } + +#define FOR_EACH_SM70_MOE_GEOMETRY(M) \ + M(32, 64, 32, 4, 32, 32, 16, 128) \ + M(32, 64, 32, 4, 32, 32, 16, 256) \ + M(32, 64, 64, 4, 32, 32, 32, 128) \ + M(32, 64, 64, 4, 32, 32, 32, 256) \ + M(32, 128, 32, 4, 32, 32, 32, 128) \ + M(32, 128, 32, 4, 32, 32, 32, 256) \ + M(32, 128, 32, 4, 32, 64, 16, 128) \ + M(32, 128, 32, 4, 32, 64, 16, 256) \ + M(32, 128, 64, 4, 32, 64, 32, 128) \ + M(32, 128, 64, 4, 32, 64, 32, 256) \ + M(32, 128, 64, 8, 32, 32, 32, 128) \ + M(32, 128, 64, 8, 32, 32, 32, 256) \ + M(32, 128, 64, 8, 32, 64, 16, 128) \ + M(32, 128, 64, 8, 32, 64, 16, 256) \ + M(32, 128, 128, 8, 32, 64, 32, 128) \ + M(32, 128, 128, 8, 32, 64, 32, 256) \ + M(32, 256, 32, 4, 32, 64, 32, 256) \ + M(32, 256, 64, 8, 32, 64, 32, 256) \ + M(64, 64, 32, 4, 32, 32, 32, 64) \ + M(64, 64, 32, 4, 32, 32, 32, 128) \ + M(64, 64, 32, 4, 32, 32, 32, 256) \ + M(64, 128, 32, 4, 32, 64, 32, 128) \ + M(64, 128, 32, 4, 32, 64, 32, 256) \ + M(64, 128, 32, 4, 64, 64, 16, 128) \ + M(64, 128, 32, 4, 64, 64, 16, 256) \ + M(64, 128, 32, 8, 32, 32, 32, 128) \ + M(64, 128, 32, 8, 32, 32, 32, 256) \ + M(64, 128, 64, 4, 64, 64, 32, 128) \ + M(64, 128, 64, 4, 64, 64, 32, 256) \ + M(64, 256, 32, 4, 64, 64, 32, 256) \ M(64, 256, 64, 8, 64, 64, 32, 256) FOR_EACH_SM70_MOE_GEOMETRY(DISPATCH_SM70_MOE_GEOMETRY) @@ -2537,8 +2507,8 @@ torch::Tensor dispatch_sm70_marlin_moe_geometry(Launcher const& launcher, int64_t group_size, char const* quant_name) { return dispatch_sm70_marlin_moe_cta_geometry( - Sm70MarlinMoeGroupSizeDispatchLauncher{ - launcher, group_size, quant_name}, + Sm70MarlinMoeGroupSizeDispatchLauncher{launcher, group_size, + quant_name}, geometry, packed_macro_n, quant_name); } @@ -2547,18 +2517,18 @@ torch::Tensor dispatch_sm70_marlin_moe_geometry_with_group32( Launcher const& launcher, Sm70CtaGeometry geometry, int packed_macro_n, int64_t group_size, char const* quant_name) { return dispatch_sm70_marlin_moe_cta_geometry( - Sm70MarlinMoeGroupSize32DispatchLauncher{ - launcher, group_size, quant_name}, + Sm70MarlinMoeGroupSize32DispatchLauncher{launcher, group_size, + quant_name}, geometry, packed_macro_n, quant_name); } template -torch::Tensor dispatch_sm70_marlin_moe_fp8_geometry( - Launcher const& launcher, Sm70CtaGeometry geometry, int packed_macro_n, - int64_t group_size) { +torch::Tensor dispatch_sm70_marlin_moe_fp8_geometry(Launcher const& launcher, + Sm70CtaGeometry geometry, + int packed_macro_n, + int64_t group_size) { return dispatch_sm70_marlin_moe_cta_geometry( - Sm70MarlinMoeFp8GroupSizeDispatchLauncher{launcher, - group_size}, + Sm70MarlinMoeFp8GroupSizeDispatchLauncher{launcher, group_size}, geometry, packed_macro_n, "fp8_e4m3"); } @@ -2603,8 +2573,8 @@ class Sm70MoeGatherIteratorA { using Element = cutlass::half_t; using Layout = cutlass::layout::RowMajor; using TensorCoord = cutlass::MatrixCoord; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + using Fragment = cutlass::Array; struct Params { int lda; int moe_block_size; @@ -2689,9 +2659,8 @@ class Sm70MoeGatherIteratorA { CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const local_row = - local_m_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const local_row = local_m_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; int const route_row = moe_block_ * params_.moe_block_size + local_row; int sorted_id = -1; @@ -2699,20 +2668,17 @@ class Sm70MoeGatherIteratorA { route_row < params_.padded_tokens; if (valid_row) { sorted_id = sorted_token_ids_[route_row]; - valid_row = - sorted_id >= 0 && sorted_id < params_.expanded_token_count; + valid_row = sorted_id >= 0 && sorted_id < params_.expanded_token_count; } int const token_row = valid_row ? (sorted_id / params_.top_k) : 0; CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const logical_k = - k_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; - int const frag_base = - (c + s * ThreadMap::Iterations::kContiguous) * - ThreadMap::kElementsPerAccess; + int const logical_k = k_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; + int const frag_base = (c + s * ThreadMap::Iterations::kContiguous) * + ThreadMap::kElementsPerAccess; bool const valid = valid_row && logical_k < params_.lda; CUTLASS_PRAGMA_UNROLL @@ -2754,8 +2720,8 @@ struct Sm70MarlinMoeGemmTraits { using LayoutB = cutlass::layout::RowMajor; using LayoutC = cutlass::layout::RowMajor; using ThreadblockShape = cutlass::gemm::GemmShape; - using WarpShape = - typename Sm70WarpShape::Type; + using WarpShape = typename Sm70WarpShape::Type; static_assert(WarpShape::kM <= 64 && WarpShape::kN <= 64, "SM70 Marlin MoE keeps per-warp M/N no larger than 64."); using InstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; @@ -2763,13 +2729,16 @@ struct Sm70MarlinMoeGemmTraits { ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, ElementB, LayoutB, ElementAccumulator, LayoutC, cutlass::arch::OpClassTensorOp, 2, cutlass::arch::OpMultiplyAdd>; - static_assert(MmaCore::kThreads == Warps * 32, - "SM70 Marlin MoE launch threads must match CUTLASS warp count."); - using IteratorA = typename Spec::template IteratorA< - ThreadblockShape, typename MmaCore::IteratorThreadMapA>; - using IteratorB = typename Spec::template IteratorB< - ThreadblockShape, typename MmaCore::IteratorThreadMapB, GroupSize, - PackedMacroN>; + static_assert( + MmaCore::kThreads == Warps * 32, + "SM70 Marlin MoE launch threads must match CUTLASS warp count."); + using IteratorA = + typename Spec::template IteratorA; + using IteratorB = + typename Spec::template IteratorB; using Mma = Sm70MarlinMmaPipelined< ThreadblockShape, IteratorA, typename MmaCore::SmemIteratorA, IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, LayoutC, @@ -2818,11 +2787,11 @@ class Sm70MoeScatterEpilogue { typename SharedLoadIterator::Fragment const& frag, int32_t const* __restrict__ sorted_token_ids, float const* __restrict__ topk_weights, - cutlass::half_t* __restrict__ c, - int n, int moe_block, int local_m_offset, - int moe_block_size, int expanded_token_count, - int padded_tokens, bool mul_topk_weights, - bool atomic_store, float output_scale) const { + cutlass::half_t* __restrict__ c, int n, int moe_block, + int local_m_offset, int moe_block_size, + int expanded_token_count, int padded_tokens, + bool mul_topk_weights, bool atomic_store, + float output_scale) const { float const* frag_ptr = reinterpret_cast(&frag); half* c_half = reinterpret_cast(c); int const thread_start_row = destination_iterator.thread_start_row(); @@ -2838,20 +2807,17 @@ class Sm70MoeScatterEpilogue { int const frag_row_idx = row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster); - int const row_offset = - row * ThreadMap::Delta::kRow + - group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; - int const local_row = - local_m_offset + thread_start_row + row_offset; + int const row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + int const local_row = local_m_offset + thread_start_row + row_offset; int const route_row = moe_block * moe_block_size + local_row; bool valid_row = local_row < moe_block_size && route_row < padded_tokens; int sorted_id = -1; if (valid_row) { sorted_id = sorted_token_ids[route_row]; - valid_row = - sorted_id >= 0 && sorted_id < expanded_token_count; + valid_row = sorted_id >= 0 && sorted_id < expanded_token_count; } float const route_scale = (valid_row && mul_topk_weights) ? topk_weights[sorted_id] : 1.0f; @@ -2896,8 +2862,7 @@ class Sm70MoeScatterEpilogue { int const warp_mn = warp_idx % (WarpCount::kM * WarpCount::kN); int const warp_m = warp_mn % WarpCount::kM; int const warp_n = warp_mn / WarpCount::kM; - cutlass::MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, - warp_n}; + cutlass::MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, warp_n}; warp_tile_iterator_.add_tile_offset(warp_offset); } @@ -2906,11 +2871,11 @@ class Sm70MoeScatterEpilogue { AccumulatorTile const& accumulators, int32_t const* __restrict__ sorted_token_ids, float const* __restrict__ topk_weights, - cutlass::half_t* __restrict__ c, - int n, int moe_block, int local_m_offset, - int moe_block_size, int expanded_token_count, - int padded_tokens, bool mul_topk_weights, - bool atomic_store, float output_scale = 1.0f) { + cutlass::half_t* __restrict__ c, int n, int moe_block, + int local_m_offset, int moe_block_size, + int expanded_token_count, int padded_tokens, + bool mul_topk_weights, bool atomic_store, + float output_scale = 1.0f) { AccumulatorFragmentIterator accum_fragment_iterator(accumulators); CUTLASS_PRAGMA_UNROLL @@ -2956,19 +2921,19 @@ class Sm70MoeScatterEpilogue { }; template -__global__ __launch_bounds__(Traits::MmaCore::kThreads, 1) -void sm70_marlin_moe_gemm_kernel( +__global__ +__launch_bounds__(Traits::MmaCore::kThreads, 1) void sm70_marlin_moe_gemm_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, typename Traits::GemmSpec::ScaleElement const* __restrict__ b_scales, typename Traits::GemmSpec::ZeroElement const* __restrict__ b_zeros, - float const* __restrict__ global_scale, - cutlass::half_t* __restrict__ c, + float const* __restrict__ global_scale, cutlass::half_t* __restrict__ c, int32_t const* __restrict__ sorted_token_ids, int32_t const* __restrict__ expert_ids, int32_t const* __restrict__ num_tokens_past_padded, float const* __restrict__ topk_weights, int moe_block_size, int top_k, - bool mul_topk_weights, int m, int n, int k, int lda, int requested_split_k) { + bool mul_topk_weights, int m, int n, int k, int lda, + int requested_split_k) { using Mma = typename Traits::Mma; using Epilogue = Sm70MoeScatterEpilogue; constexpr int CtaM = Traits::ThreadblockShape::kM; @@ -3002,8 +2967,8 @@ void sm70_marlin_moe_gemm_kernel( int partition_k = k; if constexpr (SplitK) { Sm70SplitKPartition const partition = - sm70_splitk_partition( - k, requested_split_k, int(blockIdx.z)); + sm70_splitk_partition(k, requested_split_k, + int(blockIdx.z)); if (partition.partition_k == 0) { return; } @@ -3014,8 +2979,8 @@ void sm70_marlin_moe_gemm_kernel( int const n_offset = int(blockIdx.y) * CtaN; typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(lda, moe_block_size, top_k, m, - m * top_k, padded_tokens), + typename Mma::IteratorA::Params(lda, moe_block_size, top_k, m, m * top_k, + padded_tokens), a, sorted_token_ids, thread_idx, moe_block, local_m_offset, k_begin); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), @@ -3042,8 +3007,8 @@ void sm70_marlin_moe_gemm_kernel( } Epilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx); epilogue(iterator_D, accumulators, sorted_token_ids, topk_weights, c, n, - moe_block, local_m_offset, moe_block_size, m * top_k, - padded_tokens, mul_topk_weights, SplitK, output_scale); + moe_block, local_m_offset, moe_block_size, m * top_k, padded_tokens, + mul_topk_weights, SplitK, output_scale); } template @@ -3103,9 +3068,9 @@ torch::Tensor launch_sm70_marlin_moe_gemm( smem_bytes = configure_sm70_dynamic_smem(split_kernel); int64_t const numel = size_m * top_k * size_n; - C10_CUDA_CHECK(cudaMemsetAsync( - c.data_ptr(), 0, - static_cast(numel) * sizeof(at::Half), stream)); + C10_CUDA_CHECK(cudaMemsetAsync(c.data_ptr(), 0, + static_cast(numel) * sizeof(at::Half), + stream)); int const active_split_k = sm70_active_split_k(static_cast(size_k), requested_split_k, CtaK); @@ -3115,11 +3080,11 @@ torch::Tensor launch_sm70_marlin_moe_gemm( reinterpret_cast(b_q_weight.data_ptr()), b_scales_ptr, b_zeros_ptr, global_scale_ptr, reinterpret_cast(c.data_ptr()), - sorted_token_ids.data_ptr(), - expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(), + sorted_token_ids.data_ptr(), expert_ids.data_ptr(), + num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(), int(moe_block_size), int(top_k), - mul_topk_weights, int(size_m), int(size_n), int(size_k), - int(a.stride(0)), requested_split_k); + mul_topk_weights, int(size_m), int(size_n), int(size_k), int(a.stride(0)), + requested_split_k); C10_CUDA_KERNEL_LAUNCH_CHECK(); return c; } diff --git a/csrc/moe/marlin_moe_wna16/sm70_marlin_moe_dispatch.cu b/csrc/moe/marlin_moe_wna16/sm70_marlin_moe_dispatch.cu index 733d9a4959..ec11d6d72e 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_moe_dispatch.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_moe_dispatch.cu @@ -43,15 +43,13 @@ __device__ __forceinline__ int sm70_inverse_zp_interleave(int idx, return inv[idx]; } -__device__ __forceinline__ int sm70_inverse_scale_perm(int idx, - int perm_len) { +__device__ __forceinline__ int sm70_inverse_scale_perm(int idx, int perm_len) { if (perm_len == 64) { return (idx % 8) * 8 + idx / 8; } - constexpr int inv[32] = {0, 1, 8, 9, 16, 17, 24, 25, - 2, 3, 10, 11, 18, 19, 26, 27, - 4, 5, 12, 13, 20, 21, 28, 29, - 6, 7, 14, 15, 22, 23, 30, 31}; + constexpr int inv[32] = {0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, + 11, 18, 19, 26, 27, 4, 5, 12, 13, 20, 21, + 28, 29, 6, 7, 14, 15, 22, 23, 30, 31}; return inv[idx]; } @@ -66,10 +64,9 @@ __global__ void sm70_unpermute_half_metadata_kernel( } __global__ void sm70_unpack_moe_zp_to_half_kernel( - const int32_t* __restrict__ packed_zp, - const half* __restrict__ scales, half* __restrict__ out, - int64_t total, int64_t num_groups, int64_t size_n, int64_t packed_n, - int num_bits, bool is_a_8bit) { + const int32_t* __restrict__ packed_zp, const half* __restrict__ scales, + half* __restrict__ out, int64_t total, int64_t num_groups, int64_t size_n, + int64_t packed_n, int num_bits, bool is_a_8bit) { const int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= total) return; @@ -83,15 +80,14 @@ __global__ void sm70_unpack_moe_zp_to_half_kernel( if (!is_a_8bit) { packed_col_in_block = (packed_col_in_block / pack_factor) * pack_factor + - sm70_inverse_zp_interleave(packed_col_in_block % pack_factor, - num_bits); + sm70_inverse_zp_interleave(packed_col_in_block % pack_factor, num_bits); } const int64_t packed_col = block64 * 64 + packed_col_in_block; - const uint32_t word = reinterpret_cast(packed_zp) - [(expert * num_groups + group) * packed_n + packed_col / pack_factor]; - const uint32_t zp = - (word >> ((packed_col % pack_factor) * num_bits)) & - ((1u << num_bits) - 1u); + const uint32_t word = reinterpret_cast( + packed_zp)[(expert * num_groups + group) * packed_n + + packed_col / pack_factor]; + const uint32_t zp = (word >> ((packed_col % pack_factor) * num_bits)) & + ((1u << num_bits) - 1u); out[idx] = __float2half(static_cast(zp) * __half2float(scales[idx])); } @@ -111,8 +107,8 @@ torch::Tensor sm70_moe_scales_as_logical(torch::Tensor const& b_scales, const int64_t total = out.numel(); constexpr int threads = 256; const int blocks = static_cast((total + threads - 1) / threads); - sm70_unpermute_half_metadata_kernel<<< - blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>( + sm70_unpermute_half_metadata_kernel<<>>( reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(out.data_ptr()), total, perm_len); C10_CUDA_KERNEL_LAUNCH_CHECK(); @@ -120,9 +116,8 @@ torch::Tensor sm70_moe_scales_as_logical(torch::Tensor const& b_scales, } torch::Tensor sm70_moe_zp_as_half(torch::Tensor const& b_zeros, - torch::Tensor const& b_scales, - int64_t size_n, int num_bits, - bool is_a_8bit) { + torch::Tensor const& b_scales, int64_t size_n, + int num_bits, bool is_a_8bit) { if (b_zeros.scalar_type() == at::ScalarType::Half) { return b_zeros; } @@ -143,14 +138,13 @@ torch::Tensor sm70_moe_zp_as_half(torch::Tensor const& b_zeros, TORCH_CHECK(b_scales.scalar_type() == at::ScalarType::Half, "SM70 Marlin MoE packed zero-point adapter expects fp16 scales."); - auto out = - torch::empty({b_scales.size(0), b_scales.size(1), size_n}, - b_scales.options()); + auto out = torch::empty({b_scales.size(0), b_scales.size(1), size_n}, + b_scales.options()); const int64_t total = out.numel(); constexpr int threads = 256; const int blocks = static_cast((total + threads - 1) / threads); - sm70_unpack_moe_zp_to_half_kernel<<< - blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>( + sm70_unpack_moe_zp_to_half_kernel<<>>( b_zeros.data_ptr(), reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(out.data_ptr()), total, @@ -229,8 +223,7 @@ torch::Tensor moe_wna16_marlin_gemm( std::optional const& global_scale_or_none, std::optional const& b_zeros_or_none, std::optional const& g_idx_or_none, - std::optional const& perm_or_none, - torch::Tensor& workspace, + std::optional const& perm_or_none, torch::Tensor& workspace, torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids, torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights, int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, @@ -308,8 +301,7 @@ torch::Tensor moe_wna16_marlin_gemm( "SM70 Marlin MoE supports only float16 outputs."); TORCH_CHECK(s_type == vllm::kFloat16 || (b_type == vllm::kFE2M1f && - (s_type == vllm::kFE4M3fn || - s_type == vllm::kFE8M0fnu)), + (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu)), "SM70 Marlin MoE supports only float16 scales, " "except FP4 uses float8_e4m3fn scales for NVFP4 or " "float8_e8m0fnu scales for MXFP4."); @@ -364,10 +356,8 @@ torch::Tensor moe_wna16_marlin_gemm( TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); TORCH_CHECK(b_scales.scalar_type() == at::ScalarType::Half || (b_type == vllm::kFE2M1f && - (b_scales.scalar_type() == - at::ScalarType::Float8_e4m3fn || - b_scales.scalar_type() == - at::ScalarType::Float8_e8m0fnu)), + (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn || + b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu)), "SM70 Marlin MoE supports float16 scales, except FP4 uses " "float8_e4m3fn scales for NVFP4 or float8_e8m0fnu scales for " "MXFP4."); @@ -457,17 +447,15 @@ torch::Tensor moe_wna16_marlin_gemm( torch::Tensor b_scales_kernel = b_scales; if (b_scales.scalar_type() == at::ScalarType::Half) { - b_scales_kernel = - sm70_moe_scales_as_logical(b_scales, size_k, group_size, - a_type.size_bits() == 8); + b_scales_kernel = sm70_moe_scales_as_logical(b_scales, size_k, group_size, + a_type.size_bits() == 8); } torch::Tensor global_scale; if (global_scale_or_none.has_value()) { global_scale = global_scale_or_none.value(); - TORCH_CHECK( - b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, - "SM70 Marlin MoE supports global_scale only for nvfp4 format."); + TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, + "SM70 Marlin MoE supports global_scale only for nvfp4 format."); TORCH_CHECK(global_scale.device().is_cuda(), "global_scale is not on GPU"); TORCH_CHECK(global_scale.is_contiguous(), "global_scale is not contiguous"); TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float, @@ -510,9 +498,9 @@ torch::Tensor moe_wna16_marlin_gemm( "weights when zero-points are enabled. Got = ", b_type.str()); if (!is_zp_float && b_zeros.scalar_type() == at::ScalarType::Int) { - b_zeros = sm70_moe_zp_as_half( - b_zeros, b_scales_kernel, size_n, b_type == vllm::kU4 ? 4 : 8, - a_type.size_bits() == 8); + b_zeros = sm70_moe_zp_as_half(b_zeros, b_scales_kernel, size_n, + b_type == vllm::kU4 ? 4 : 8, + a_type.size_bits() == 8); zp_is_float = true; } } else { @@ -535,8 +523,7 @@ torch::Tensor moe_wna16_marlin_gemm( TORCH_CHECK(num_groups == b_zeros.size(1), "b_zeros dim 1 = ", b_zeros.size(1), " is not num_groups = ", num_groups); - TORCH_CHECK(b_zeros.size(2) == size_n, - "b_zeros dim 2 = ", b_zeros.size(2), + TORCH_CHECK(b_zeros.size(2) == size_n, "b_zeros dim 2 = ", b_zeros.size(2), " is not size_n = ", size_n); } else { TORCH_CHECK(!zp_is_float, @@ -559,20 +546,18 @@ torch::Tensor moe_wna16_marlin_gemm( TORCH_CHECK(!has_act_order, "act_order is not supported for the SM70 Marlin MoE path."); - TORCH_CHECK(!has_bias, - "SM70 Marlin MoE does not support bias."); + TORCH_CHECK(!has_bias, "SM70 Marlin MoE does not support bias."); TORCH_CHECK(!use_atomic_add, "SM70 Marlin MoE does not support atomic-add " "output."); - TORCH_CHECK(is_k_full, - "SM70 Marlin MoE requires full-K inputs."); + TORCH_CHECK(is_k_full, "SM70 Marlin MoE requires full-K inputs."); if (b_type == vllm::kU4) { TORCH_CHECK(has_zp && zp_is_float, "SM70 Marlin MoE uint4 path requires fp16 zero points."); return MARLIN_NAMESPACE_NAME::sm70_marlin_u4_gemm( - a, c, b_q_weight, b_scales_kernel, b_zeros, sorted_token_ids, expert_ids, - num_tokens_past_padded, topk_weights, moe_block_size, top_k, + a, c, b_q_weight, b_scales_kernel, b_zeros, sorted_token_ids, + expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, mul_topk_weights, size_m, size_n, size_k, group_size); } @@ -580,8 +565,8 @@ torch::Tensor moe_wna16_marlin_gemm( TORCH_CHECK(has_zp && zp_is_float, "SM70 Marlin MoE uint8 path requires fp16 zero points."); return MARLIN_NAMESPACE_NAME::sm70_marlin_u8_gemm( - a, c, b_q_weight, b_scales_kernel, b_zeros, sorted_token_ids, expert_ids, - num_tokens_past_padded, topk_weights, moe_block_size, top_k, + a, c, b_q_weight, b_scales_kernel, b_zeros, sorted_token_ids, + expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, mul_topk_weights, size_m, size_n, size_k, group_size); } @@ -622,8 +607,9 @@ torch::Tensor moe_wna16_marlin_gemm( TORCH_CHECK(!has_zp && !is_zp_float, "SM70 Marlin MoE fp4 does not support zero-point metadata."); if (s_type == vllm::kFE4M3fn) { - TORCH_CHECK(global_scale.numel() == num_experts, - "the global_scale parameter must be passed for nvfp4 format."); + TORCH_CHECK( + global_scale.numel() == num_experts, + "the global_scale parameter must be passed for nvfp4 format."); TORCH_CHECK(group_size == 16, "SM70 Marlin MoE NVFP4 supports only group_size 16. " "Got ", diff --git a/csrc/moe/marlin_moe_wna16/sm70_marlin_mxfp4_gemm.cu b/csrc/moe/marlin_moe_wna16/sm70_marlin_mxfp4_gemm.cu index 76443f12b8..9755b3edbf 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_mxfp4_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_mxfp4_gemm.cu @@ -51,7 +51,8 @@ void dequant_e8m0_scales_to_half2(int q, half2* scale_cache) { *reinterpret_cast(scale_cache) = vec_res; } -template +template class Sm70MoeMxfp4IteratorB { public: using Shape = Shape_; @@ -59,22 +60,25 @@ class Sm70MoeMxfp4IteratorB { static int const kGroupSize = GroupSize_; static int const kPackedMacroN = PackedMacroN_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; - static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, - "SM70 Marlin MoE MXFP4 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, - "SM70 Marlin MoE MXFP4 IteratorB expects one contiguous iteration " - "per 64-column quant tile."); + using Fragment = cutlass::Array; + static_assert( + Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, + "SM70 Marlin MoE MXFP4 IteratorB expects CTA_N in {64, 128, 256}."); + static_assert( + ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, + "SM70 Marlin MoE MXFP4 IteratorB expects one contiguous iteration " + "per 64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, "SM70 Marlin MoE MXFP4 IteratorB expects 64-column deltas."); - static_assert(ThreadMap::kElementsPerAccess == kMxfp4ValuesPerWord, - "SM70 Marlin MoE MXFP4 IteratorB expects one packed FP4 word per " - "access."); - static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin MoE U4-family IteratorB expects one or more strided " - "iterations."); + static_assert( + ThreadMap::kElementsPerAccess == kMxfp4ValuesPerWord, + "SM70 Marlin MoE MXFP4 IteratorB expects one packed FP4 word per " + "access."); + static_assert( + ThreadMap::Iterations::kStrided >= 1, + "SM70 Marlin MoE U4-family IteratorB expects one or more strided " + "iterations."); struct Params { int size_k; int size_n; @@ -85,9 +89,7 @@ class Sm70MoeMxfp4IteratorB { CUTLASS_HOST_DEVICE Params(int size_k_, int size_n_) - : size_k(size_k_), - size_n(size_n_), - num_groups(size_k_ / GroupSize_) {} + : size_k(size_k_), size_n(size_n_), num_groups(size_k_ / GroupSize_) {} }; private: @@ -108,9 +110,8 @@ class Sm70MoeMxfp4IteratorB { uint8_t const* scales, half const*, int thread_id, int expert, int k_offset, int n_offset) : qweight_(qweight), - scales_expert_base_(scales + - (expert >= 0 ? expert : 0) * - params.num_groups * params.size_n), + scales_expert_base_(scales + (expert >= 0 ? expert : 0) * + params.num_groups * params.size_n), params_(params), thread_offset_(ThreadMap::initial_offset(thread_id)), expert_(expert), @@ -149,9 +150,7 @@ class Sm70MoeMxfp4IteratorB { } CUTLASS_DEVICE - int metadata_group_offset(int group) const { - return group * params_.size_n; - } + int metadata_group_offset(int group) const { return group * params_.size_n; } CUTLASS_DEVICE int scale_group(int logical_k) const { @@ -184,14 +183,12 @@ class Sm70MoeMxfp4IteratorB { CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { int const logical_k = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + k_offset_ + thread_offset_.strided() + s * ThreadMap::Delta::kStrided; int const group = scale_group(logical_k); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; int const cache_index = c + s * ThreadMap::Iterations::kContiguous; cache_metadata_e8m0_scales(cache_index, group, cache_n); } @@ -244,8 +241,9 @@ class Sm70MoeMxfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE MXFP4 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE MXFP4 contiguous iteration count."); uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_offset_); half2 const* scale_vec = cached_scales_; @@ -265,15 +263,13 @@ class Sm70MoeMxfp4IteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; int const qweight_base_s = expert_qweight_base_offset() + qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - auto const qwords = - load_qword_vector<4>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<4>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -296,8 +292,7 @@ class Sm70MoeMxfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - auto const qwords = - load_qword_vector<2>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<2>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -320,8 +315,9 @@ class Sm70MoeMxfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE MXFP4 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE MXFP4 contiguous iteration count."); uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_s); constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -364,15 +360,15 @@ struct Sm70MoeMxfp4GemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = Sm70MoeMxfp4IteratorB; + using IteratorB = + Sm70MoeMxfp4IteratorB; }; -template +template using Sm70MoeMxfp4GemmTraits = - Sm70MarlinMoeGemmTraits; + Sm70MarlinMoeGemmTraits; struct Sm70MoeMxfp4Launcher { torch::Tensor& a; @@ -393,15 +389,15 @@ struct Sm70MoeMxfp4Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - using Traits = Sm70MoeMxfp4GemmTraits; + using Traits = Sm70MoeMxfp4GemmTraits; return launch_sm70_marlin_moe_gemm( a, c, b_q_weight, b_scales, b_zeros, global_scale, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, - top_k, mul_topk_weights, size_m, size_n, size_k, requested_split_k); + expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, + mul_topk_weights, size_m, size_n, size_k, requested_split_k); } }; @@ -419,9 +415,10 @@ torch::Tensor sm70_marlin_mxfp4_gemm( auto const params = sm70_marlin_moe_auto_stage_params( "mxfp4", group_size, moe_block_size, top_k, size_m, size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE mxfp4", geometry); - validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE mxfp4", geometry, - size_n); + validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE mxfp4", + geometry); + validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE mxfp4", + geometry, size_n); TORCH_CHECK(size_k % geometry.cta_k == 0, "SM70 Marlin MoE mxfp4 requires K divisible by CTA_K=", geometry.cta_k, ". Got K=", size_k, "."); @@ -429,10 +426,23 @@ torch::Tensor sm70_marlin_mxfp4_gemm( auto empty_half = torch::empty({0}, b_scales.options().dtype(at::kHalf)); auto empty_float = torch::empty( {0}, torch::TensorOptions().dtype(at::kFloat).device(a.device())); - Sm70MoeMxfp4Launcher const launcher{ - a, c, b_q_weight, b_scales, empty_half, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; + Sm70MoeMxfp4Launcher const launcher{a, + c, + b_q_weight, + b_scales, + empty_half, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; return dispatch_sm70_marlin_moe_fixed_group_geometry<32>( launcher, geometry, params.packed_macro_n, group_size, "mxfp4"); } diff --git a/csrc/moe/marlin_moe_wna16/sm70_marlin_nvfp4_gemm.cu b/csrc/moe/marlin_moe_wna16/sm70_marlin_nvfp4_gemm.cu index a1ea955d34..b2cc86aa91 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_nvfp4_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_nvfp4_gemm.cu @@ -26,7 +26,8 @@ namespace { constexpr int kNvfp4ValuesPerWord = 8; -template +template class Sm70MoeNvfp4IteratorB { public: using Shape = Shape_; @@ -34,22 +35,25 @@ class Sm70MoeNvfp4IteratorB { static int const kGroupSize = GroupSize_; static int const kPackedMacroN = PackedMacroN_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; - static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, - "SM70 Marlin MoE NVFP4 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, - "SM70 Marlin MoE NVFP4 IteratorB expects one contiguous iteration " - "per 64-column quant tile."); + using Fragment = cutlass::Array; + static_assert( + Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, + "SM70 Marlin MoE NVFP4 IteratorB expects CTA_N in {64, 128, 256}."); + static_assert( + ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, + "SM70 Marlin MoE NVFP4 IteratorB expects one contiguous iteration " + "per 64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, "SM70 Marlin MoE NVFP4 IteratorB expects 64-column deltas."); - static_assert(ThreadMap::kElementsPerAccess == kNvfp4ValuesPerWord, - "SM70 Marlin MoE NVFP4 IteratorB expects one packed FP4 word per " - "access."); - static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin MoE U4-family IteratorB expects one or more strided " - "iterations."); + static_assert( + ThreadMap::kElementsPerAccess == kNvfp4ValuesPerWord, + "SM70 Marlin MoE NVFP4 IteratorB expects one packed FP4 word per " + "access."); + static_assert( + ThreadMap::Iterations::kStrided >= 1, + "SM70 Marlin MoE U4-family IteratorB expects one or more strided " + "iterations."); struct Params { int size_k; int size_n; @@ -60,9 +64,7 @@ class Sm70MoeNvfp4IteratorB { CUTLASS_HOST_DEVICE Params(int size_k_, int size_n_) - : size_k(size_k_), - size_n(size_n_), - num_groups(size_k_ / GroupSize_) {} + : size_k(size_k_), size_n(size_n_), num_groups(size_k_ / GroupSize_) {} }; private: @@ -83,9 +85,8 @@ class Sm70MoeNvfp4IteratorB { uint8_t const* scales, half const*, int thread_id, int expert, int k_offset, int n_offset) : qweight_(qweight), - scales_expert_base_(scales + - (expert >= 0 ? expert : 0) * - params.num_groups * params.size_n), + scales_expert_base_(scales + (expert >= 0 ? expert : 0) * + params.num_groups * params.size_n), params_(params), thread_offset_(ThreadMap::initial_offset(thread_id)), expert_(expert), @@ -124,9 +125,7 @@ class Sm70MoeNvfp4IteratorB { } CUTLASS_DEVICE - int metadata_group_offset(int group) const { - return group * params_.size_n; - } + int metadata_group_offset(int group) const { return group * params_.size_n; } CUTLASS_DEVICE int scale_group(int logical_k) const { @@ -159,14 +158,12 @@ class Sm70MoeNvfp4IteratorB { CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { int const logical_k = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + k_offset_ + thread_offset_.strided() + s * ThreadMap::Delta::kStrided; int const group = scale_group(logical_k); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; int const cache_index = c + s * ThreadMap::Iterations::kContiguous; cache_metadata_fp8_scales(cache_index, group, cache_n); } @@ -219,8 +216,9 @@ class Sm70MoeNvfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE NVFP4 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE NVFP4 contiguous iteration count."); uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_offset_); half2 const* scale_vec = cached_scales_; @@ -240,15 +238,13 @@ class Sm70MoeNvfp4IteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; int const qweight_base_s = expert_qweight_base_offset() + qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - auto const qwords = - load_qword_vector<4>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<4>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -271,8 +267,7 @@ class Sm70MoeNvfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - auto const qwords = - load_qword_vector<2>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<2>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -295,8 +290,9 @@ class Sm70MoeNvfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE NVFP4 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE NVFP4 contiguous iteration count."); uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_s); constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -339,15 +335,15 @@ struct Sm70MoeNvfp4GemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = Sm70MoeNvfp4IteratorB; + using IteratorB = + Sm70MoeNvfp4IteratorB; }; -template +template using Sm70MoeNvfp4GemmTraits = - Sm70MarlinMoeGemmTraits; + Sm70MarlinMoeGemmTraits; struct Sm70MoeNvfp4Launcher { torch::Tensor& a; @@ -368,15 +364,15 @@ struct Sm70MoeNvfp4Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - using Traits = Sm70MoeNvfp4GemmTraits; + using Traits = Sm70MoeNvfp4GemmTraits; return launch_sm70_marlin_moe_gemm( a, c, b_q_weight, b_scales, b_zeros, global_scale, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, - top_k, mul_topk_weights, size_m, size_n, size_k, requested_split_k); + expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, + mul_topk_weights, size_m, size_n, size_k, requested_split_k); } }; @@ -394,18 +390,32 @@ torch::Tensor sm70_marlin_nvfp4_gemm( auto const params = sm70_marlin_moe_auto_stage_params( "nvfp4", group_size, moe_block_size, top_k, size_m, size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE nvfp4", geometry); - validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE nvfp4", geometry, - size_n); + validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE nvfp4", + geometry); + validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE nvfp4", + geometry, size_n); TORCH_CHECK(size_k % geometry.cta_k == 0, "SM70 Marlin MoE nvfp4 requires K divisible by CTA_K=", geometry.cta_k, ". Got K=", size_k, "."); auto empty_half = torch::empty({0}, b_scales.options().dtype(at::kHalf)); - Sm70MoeNvfp4Launcher const launcher{ - a, c, b_q_weight, b_scales, empty_half, global_scale, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; + Sm70MoeNvfp4Launcher const launcher{a, + c, + b_q_weight, + b_scales, + empty_half, + global_scale, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; return dispatch_sm70_marlin_moe_fixed_group_geometry<16>( launcher, geometry, params.packed_macro_n, group_size, "nvfp4"); } diff --git a/csrc/moe/marlin_moe_wna16/sm70_marlin_u4_gemm.cu b/csrc/moe/marlin_moe_wna16/sm70_marlin_u4_gemm.cu index cf412f4d0f..bce2bfb899 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_u4_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_u4_gemm.cu @@ -35,21 +35,23 @@ class Sm70MoeU4ZpIteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; - static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, - "SM70 Marlin MoE U4 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, - "SM70 Marlin MoE U4 IteratorB expects one contiguous iteration per " - "64-column quant tile."); + using Fragment = cutlass::Array; + static_assert( + Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, + "SM70 Marlin MoE U4 IteratorB expects CTA_N in {64, 128, 256}."); + static_assert( + ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, + "SM70 Marlin MoE U4 IteratorB expects one contiguous iteration per " + "64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, "SM70 Marlin MoE U4 IteratorB expects 64-column deltas."); static_assert(ThreadMap::kElementsPerAccess == kU4ValuesPerWord, "SM70 Marlin MoE U4 IteratorB expects one packed int4 word per " "access."); static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin MoE U4-family IteratorB expects one or more strided iterations."); + "SM70 Marlin MoE U4-family IteratorB expects one or more " + "strided iterations."); struct Params { int size_k; int size_n; @@ -85,11 +87,10 @@ class Sm70MoeU4ZpIteratorB { half const* scales, half const* zp, int thread_id, int expert, int k_offset, int n_offset) : qweight_(qweight), - scales_expert_base_(scales + - (expert >= 0 ? expert : 0) * - params.num_groups * params.size_n), - zp_expert_base_(zp + (expert >= 0 ? expert : 0) * - params.num_groups * params.size_n), + scales_expert_base_(scales + (expert >= 0 ? expert : 0) * + params.num_groups * params.size_n), + zp_expert_base_(zp + (expert >= 0 ? expert : 0) * params.num_groups * + params.size_n), params_(params), thread_offset_(ThreadMap::initial_offset(thread_id)), expert_(expert), @@ -131,17 +132,14 @@ class Sm70MoeU4ZpIteratorB { } CUTLASS_DEVICE - int metadata_group_offset(int group) const { - return group * params_.size_n; - } + int metadata_group_offset(int group) const { return group * params_.size_n; } CUTLASS_DEVICE int scale_group(int logical_k) const { if constexpr (kGroupSize == -1) { return 0; } else { - static_assert(kGroupSize == 32 || kGroupSize == 64 || - kGroupSize == 128, + static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, "SM70 Marlin MoE U4 supports only group sizes -1, 32, 64, " "and 128."); return logical_k / kGroupSize; @@ -158,16 +156,16 @@ class Sm70MoeU4ZpIteratorB { CUTLASS_DEVICE void cache_metadata_lane_vectors(int c, int group, int cache_n) const { int const metadata_offset = metadata_group_offset(group) + cache_n; - half2 const* scale_vec = reinterpret_cast( - scales_expert_base_ + metadata_offset); + half2 const* scale_vec = + reinterpret_cast(scales_expert_base_ + metadata_offset); half2* scale_cache = cached_scales_ + c * 4; scale_cache[0] = scale_vec[0]; scale_cache[1] = scale_vec[1]; scale_cache[2] = scale_vec[2]; scale_cache[3] = scale_vec[3]; - half2 const* zp_vec = reinterpret_cast( - zp_expert_base_ + metadata_offset); + half2 const* zp_vec = + reinterpret_cast(zp_expert_base_ + metadata_offset); half2* zp_cache = cached_zp_ + c * 4; zp_cache[0] = zp_vec[0]; zp_cache[1] = zp_vec[1]; @@ -178,8 +176,8 @@ class Sm70MoeU4ZpIteratorB { CUTLASS_DEVICE void cache_metadata_vector_words(int c, int group, int cache_n) const { int const metadata_offset = metadata_group_offset(group) + cache_n; - uint4 const scale_words = *reinterpret_cast( - scales_expert_base_ + metadata_offset); + uint4 const scale_words = + *reinterpret_cast(scales_expert_base_ + metadata_offset); half2 const* scale_vec = reinterpret_cast(&scale_words); half2* scale_cache = cached_scales_ + c * 4; scale_cache[0] = scale_vec[0]; @@ -187,8 +185,8 @@ class Sm70MoeU4ZpIteratorB { scale_cache[2] = scale_vec[2]; scale_cache[3] = scale_vec[3]; - uint4 const zp_words = *reinterpret_cast( - zp_expert_base_ + metadata_offset); + uint4 const zp_words = + *reinterpret_cast(zp_expert_base_ + metadata_offset); half2 const* zp_vec = reinterpret_cast(&zp_words); half2* zp_cache = cached_zp_ + c * 4; zp_cache[0] = zp_vec[0]; @@ -201,9 +199,8 @@ class Sm70MoeU4ZpIteratorB { void cache_current_group_metadata(int group) const { CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; if constexpr (kUseMetadataVectorWords) { cache_metadata_vector_words(c, group, cache_n); } else { @@ -228,8 +225,8 @@ class Sm70MoeU4ZpIteratorB { half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data() + frag_base); - marlin::dequant( - static_cast(qword), deq); + marlin::dequant(static_cast(qword), + deq); frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( @@ -250,8 +247,8 @@ class Sm70MoeU4ZpIteratorB { half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data() + frag_base); - marlin::dequant( - static_cast(qword), deq); + marlin::dequant(static_cast(qword), + deq); frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( @@ -260,16 +257,18 @@ class Sm70MoeU4ZpIteratorB { frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE U4 contiguous iteration count."); - uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_offset_); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE U4 contiguous iteration count."); + uint32_t const qword = + load_qword_vector<1>(qweight_ + qweight_base_offset_); half2 const* scale_vec = cached_scales_; half2 const* zp_vec = cached_zp_; half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data()); - marlin::dequant( - static_cast(qword), deq); + marlin::dequant(static_cast(qword), + deq); frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( @@ -281,9 +280,8 @@ class Sm70MoeU4ZpIteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; if constexpr (kGroupSize != -1) { cache_current_group_metadata(scale_group(logical_k_s)); } @@ -291,8 +289,7 @@ class Sm70MoeU4ZpIteratorB { expert_qweight_base_offset() + qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - auto const qwords = - load_qword_vector<4>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<4>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -314,8 +311,7 @@ class Sm70MoeU4ZpIteratorB { frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - auto const qwords = - load_qword_vector<2>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<2>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -337,9 +333,11 @@ class Sm70MoeU4ZpIteratorB { frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE U4 contiguous iteration count."); - uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_s); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE U4 contiguous iteration count."); + uint32_t const qword = + load_qword_vector<1>(qweight_ + qweight_base_s); constexpr int kAccess = ThreadMap::kElementsPerAccess; int const frag_base = s * kAccess; half2 const* scale_vec = cached_scales_; @@ -347,8 +345,8 @@ class Sm70MoeU4ZpIteratorB { half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data() + frag_base); - marlin::dequant( - static_cast(qword), deq); + marlin::dequant(static_cast(qword), + deq); frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( @@ -366,8 +364,7 @@ class Sm70MoeU4ZpIteratorB { return; } - if constexpr (kGroupSize != -1 && - ThreadMap::Iterations::kStrided == 1) { + if constexpr (kGroupSize != -1 && ThreadMap::Iterations::kStrided == 1) { int const first_logical_k = k_offset_ + thread_offset_.strided(); cache_current_group_metadata(scale_group(first_logical_k)); } @@ -386,18 +383,17 @@ struct Sm70MoeU4ZpGemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = - Sm70MoeU4ZpIteratorB; + using IteratorB = Sm70MoeU4ZpIteratorB; }; -template using Sm70MoeU4ZpGemmTraits = - Sm70MarlinMoeGemmTraits< - Sm70MoeU4ZpGemmSpec, CtaM, CtaN, CtaK, - Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; + Sm70MarlinMoeGemmTraits, CtaM, + CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>; template struct Sm70MoeU4Launcher { @@ -419,16 +415,16 @@ struct Sm70MoeU4Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - using Traits = Sm70MoeU4ZpGemmTraits; + using Traits = + Sm70MoeU4ZpGemmTraits; return launch_sm70_marlin_moe_gemm( a, c, b_q_weight, b_scales, b_zeros, global_scale, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, - top_k, mul_topk_weights, size_m, size_n, size_k, requested_split_k); + expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, + mul_topk_weights, size_m, size_n, size_k, requested_split_k); } }; @@ -446,9 +442,10 @@ torch::Tensor sm70_marlin_u4_gemm( auto const params = sm70_marlin_moe_auto_stage_params( "uint4", group_size, moe_block_size, top_k, size_m, size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE uint4", geometry); - validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE uint4", geometry, - size_n); + validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE uint4", + geometry); + validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE uint4", + geometry, size_n); TORCH_CHECK(size_k % geometry.cta_k == 0, "SM70 Marlin MoE uint4 requires K divisible by CTA_K=", geometry.cta_k, ". Got K=", size_k, "."); @@ -456,19 +453,45 @@ torch::Tensor sm70_marlin_u4_gemm( auto empty_float = torch::empty( {0}, torch::TensorOptions().dtype(at::kFloat).device(a.device())); if (params.use_metadata_vector_words) { - Sm70MoeU4Launcher const launcher{ - a, c, b_q_weight, b_scales, b_zeros, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; - return dispatch_sm70_marlin_moe_geometry(launcher, geometry, params.packed_macro_n, group_size, - "uint4"); + Sm70MoeU4Launcher const launcher{a, + c, + b_q_weight, + b_scales, + b_zeros, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; + return dispatch_sm70_marlin_moe_geometry( + launcher, geometry, params.packed_macro_n, group_size, "uint4"); } - Sm70MoeU4Launcher const launcher{ - a, c, b_q_weight, b_scales, b_zeros, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; - return dispatch_sm70_marlin_moe_geometry(launcher, geometry, params.packed_macro_n, group_size, - "uint4"); + Sm70MoeU4Launcher const launcher{a, + c, + b_q_weight, + b_scales, + b_zeros, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; + return dispatch_sm70_marlin_moe_geometry( + launcher, geometry, params.packed_macro_n, group_size, "uint4"); } } // namespace marlin_moe_wna16 diff --git a/csrc/moe/marlin_moe_wna16/sm70_marlin_u4b8_gemm.cu b/csrc/moe/marlin_moe_wna16/sm70_marlin_u4b8_gemm.cu index 56f4206460..c5e03f0ece 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_u4b8_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_u4b8_gemm.cu @@ -35,22 +35,25 @@ class Sm70MoeU4B8IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; - static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, - "SM70 Marlin MoE U4B8 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, - "SM70 Marlin MoE U4B8 IteratorB expects one contiguous iteration per " - "64-column quant tile."); + using Fragment = cutlass::Array; + static_assert( + Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, + "SM70 Marlin MoE U4B8 IteratorB expects CTA_N in {64, 128, 256}."); + static_assert( + ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, + "SM70 Marlin MoE U4B8 IteratorB expects one contiguous iteration per " + "64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, "SM70 Marlin MoE U4B8 IteratorB expects 64-column deltas."); - static_assert(ThreadMap::kElementsPerAccess == kU4ValuesPerWord, - "SM70 Marlin MoE U4B8 IteratorB expects one packed int4 word per " - "access."); - static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin MoE U4-family IteratorB expects one or more strided " - "iterations."); + static_assert( + ThreadMap::kElementsPerAccess == kU4ValuesPerWord, + "SM70 Marlin MoE U4B8 IteratorB expects one packed int4 word per " + "access."); + static_assert( + ThreadMap::Iterations::kStrided >= 1, + "SM70 Marlin MoE U4-family IteratorB expects one or more strided " + "iterations."); struct Params { int size_k; int size_n; @@ -84,9 +87,8 @@ class Sm70MoeU4B8IteratorB { half const* scales, half const*, int thread_id, int expert, int k_offset, int n_offset) : qweight_(qweight), - scales_expert_base_(scales + - (expert >= 0 ? expert : 0) * - params.num_groups * params.size_n), + scales_expert_base_(scales + (expert >= 0 ? expert : 0) * + params.num_groups * params.size_n), params_(params), thread_offset_(ThreadMap::initial_offset(thread_id)), expert_(expert), @@ -128,19 +130,17 @@ class Sm70MoeU4B8IteratorB { } CUTLASS_DEVICE - int metadata_group_offset(int group) const { - return group * params_.size_n; - } + int metadata_group_offset(int group) const { return group * params_.size_n; } CUTLASS_DEVICE int scale_group(int logical_k) const { if constexpr (kGroupSize == -1) { return 0; } else { - static_assert(kGroupSize == 32 || kGroupSize == 64 || - kGroupSize == 128, - "SM70 Marlin MoE U4B8 supports only group sizes -1, 32, 64, " - "and 128."); + static_assert( + kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, + "SM70 Marlin MoE U4B8 supports only group sizes -1, 32, 64, " + "and 128."); return logical_k / kGroupSize; } } @@ -155,8 +155,8 @@ class Sm70MoeU4B8IteratorB { CUTLASS_DEVICE void cache_metadata_lane_vectors(int c, int group, int cache_n) const { int const metadata_offset = metadata_group_offset(group) + cache_n; - half2 const* scale_vec = reinterpret_cast( - scales_expert_base_ + metadata_offset); + half2 const* scale_vec = + reinterpret_cast(scales_expert_base_ + metadata_offset); half2* scale_cache = cached_scales_ + c * 4; scale_cache[0] = scale_vec[0]; scale_cache[1] = scale_vec[1]; @@ -167,8 +167,8 @@ class Sm70MoeU4B8IteratorB { CUTLASS_DEVICE void cache_metadata_vector_words(int c, int group, int cache_n) const { int const metadata_offset = metadata_group_offset(group) + cache_n; - uint4 const scale_words = *reinterpret_cast( - scales_expert_base_ + metadata_offset); + uint4 const scale_words = + *reinterpret_cast(scales_expert_base_ + metadata_offset); half2 const* scale_vec = reinterpret_cast(&scale_words); half2* scale_cache = cached_scales_ + c * 4; scale_cache[0] = scale_vec[0]; @@ -181,9 +181,8 @@ class Sm70MoeU4B8IteratorB { void cache_current_group_metadata(int group) const { CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; if constexpr (kUseMetadataVectorWords) { cache_metadata_vector_words(c, group, cache_n); } else { @@ -238,16 +237,17 @@ class Sm70MoeU4B8IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE U4B8 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE U4B8 contiguous iteration count."); uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_offset_); half2 const* scale_vec = cached_scales_; half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data()); - marlin::dequant( - static_cast(qword), deq); + marlin::dequant(static_cast(qword), + deq); frag_vec[0] = __hmul2(deq[0], scale_vec[0]); frag_vec[1] = __hmul2(deq[1], scale_vec[1]); marlin::dequant( @@ -259,9 +259,8 @@ class Sm70MoeU4B8IteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; if constexpr (kGroupSize != -1) { cache_current_group_metadata(scale_group(logical_k_s)); } @@ -269,8 +268,7 @@ class Sm70MoeU4B8IteratorB { expert_qweight_base_offset() + qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - auto const qwords = - load_qword_vector<4>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<4>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -291,8 +289,7 @@ class Sm70MoeU4B8IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - auto const qwords = - load_qword_vector<2>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<2>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -313,8 +310,9 @@ class Sm70MoeU4B8IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE U4B8 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE U4B8 contiguous iteration count."); uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_s); constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -342,8 +340,7 @@ class Sm70MoeU4B8IteratorB { return; } - if constexpr (kGroupSize != -1 && - ThreadMap::Iterations::kStrided == 1) { + if constexpr (kGroupSize != -1 && ThreadMap::Iterations::kStrided == 1) { int const first_logical_k = k_offset_ + thread_offset_.strided(); cache_current_group_metadata(scale_group(first_logical_k)); } @@ -362,18 +359,17 @@ struct Sm70MoeU4B8GemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = - Sm70MoeU4B8IteratorB; + using IteratorB = Sm70MoeU4B8IteratorB; }; -template using Sm70MoeU4B8GemmTraits = - Sm70MarlinMoeGemmTraits< - Sm70MoeU4B8GemmSpec, CtaM, CtaN, CtaK, - Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; + Sm70MarlinMoeGemmTraits, CtaM, + CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>; template struct Sm70MoeU4B8Launcher { @@ -395,16 +391,16 @@ struct Sm70MoeU4B8Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - using Traits = Sm70MoeU4B8GemmTraits; + using Traits = + Sm70MoeU4B8GemmTraits; return launch_sm70_marlin_moe_gemm( a, c, b_q_weight, b_scales, b_zeros, global_scale, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, - top_k, mul_topk_weights, size_m, size_n, size_k, requested_split_k); + expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, + mul_topk_weights, size_m, size_n, size_k, requested_split_k); } }; @@ -422,9 +418,10 @@ torch::Tensor sm70_marlin_u4b8_gemm( auto const params = sm70_marlin_moe_auto_stage_params( "uint4b8", group_size, moe_block_size, top_k, size_m, size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE uint4b8", geometry); - validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE uint4b8", geometry, - size_n); + validate_sm70_marlin_moe_stage_cta_geometry_supported( + "SM70 Marlin MoE uint4b8", geometry); + validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE uint4b8", + geometry, size_n); TORCH_CHECK(size_k % geometry.cta_k == 0, "SM70 Marlin MoE uint4b8 requires K divisible by CTA_K=", geometry.cta_k, ". Got K=", size_k, "."); @@ -433,17 +430,43 @@ torch::Tensor sm70_marlin_u4b8_gemm( auto empty_float = torch::empty( {0}, torch::TensorOptions().dtype(at::kFloat).device(a.device())); if (params.use_metadata_vector_words) { - Sm70MoeU4B8Launcher const launcher{ - a, c, b_q_weight, b_scales, empty_half, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; + Sm70MoeU4B8Launcher const launcher{a, + c, + b_q_weight, + b_scales, + empty_half, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; return dispatch_sm70_marlin_moe_geometry( launcher, geometry, params.packed_macro_n, group_size, "uint4b8"); } - Sm70MoeU4B8Launcher const launcher{ - a, c, b_q_weight, b_scales, empty_half, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; + Sm70MoeU4B8Launcher const launcher{a, + c, + b_q_weight, + b_scales, + empty_half, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; return dispatch_sm70_marlin_moe_geometry( launcher, geometry, params.packed_macro_n, group_size, "uint4b8"); } diff --git a/csrc/moe/marlin_moe_wna16/sm70_marlin_u8_gemm.cu b/csrc/moe/marlin_moe_wna16/sm70_marlin_u8_gemm.cu index c620f30112..25c86f4f2a 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_u8_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_u8_gemm.cu @@ -35,22 +35,24 @@ class Sm70MoeU8ZpIteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; - static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, - "SM70 Marlin MoE U8 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, - "SM70 Marlin MoE U8 IteratorB expects one contiguous iteration per " - "64-column quant tile."); + using Fragment = cutlass::Array; + static_assert( + Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, + "SM70 Marlin MoE U8 IteratorB expects CTA_N in {64, 128, 256}."); + static_assert( + ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, + "SM70 Marlin MoE U8 IteratorB expects one contiguous iteration per " + "64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, "SM70 Marlin MoE U8 IteratorB expects 64-column deltas."); static_assert(ThreadMap::kElementsPerAccess == kU8ValuesPerAccess, "SM70 Marlin MoE U8 IteratorB expects two packed U8 words per " "access."); - static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin MoE U8-family IteratorB expects one or more strided " - "iterations."); + static_assert( + ThreadMap::Iterations::kStrided >= 1, + "SM70 Marlin MoE U8-family IteratorB expects one or more strided " + "iterations."); static constexpr int kQweightWordStrideWords = kPackedMacroN / kQuantTileN; struct Params { @@ -88,11 +90,10 @@ class Sm70MoeU8ZpIteratorB { half const* scales, half const* zp, int thread_id, int expert, int k_offset, int n_offset) : qweight_(qweight), - scales_expert_base_(scales + - (expert >= 0 ? expert : 0) * - params.num_groups * params.size_n), - zp_expert_base_(zp + (expert >= 0 ? expert : 0) * - params.num_groups * params.size_n), + scales_expert_base_(scales + (expert >= 0 ? expert : 0) * + params.num_groups * params.size_n), + zp_expert_base_(zp + (expert >= 0 ? expert : 0) * params.num_groups * + params.size_n), params_(params), thread_offset_(ThreadMap::initial_offset(thread_id)), expert_(expert), @@ -134,17 +135,14 @@ class Sm70MoeU8ZpIteratorB { } CUTLASS_DEVICE - int metadata_group_offset(int group) const { - return group * params_.size_n; - } + int metadata_group_offset(int group) const { return group * params_.size_n; } CUTLASS_DEVICE int scale_group(int logical_k) const { if constexpr (kGroupSize == -1) { return 0; } else { - static_assert(kGroupSize == 32 || kGroupSize == 64 || - kGroupSize == 128, + static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, "SM70 Marlin MoE U8 supports only group sizes -1, 32, 64, " "and 128."); return logical_k / kGroupSize; @@ -161,16 +159,16 @@ class Sm70MoeU8ZpIteratorB { CUTLASS_DEVICE void cache_metadata_lane_vectors(int c, int group, int cache_n) const { int const metadata_offset = metadata_group_offset(group) + cache_n; - half2 const* scale_vec = reinterpret_cast( - scales_expert_base_ + metadata_offset); + half2 const* scale_vec = + reinterpret_cast(scales_expert_base_ + metadata_offset); half2* scale_cache = cached_scales_ + c * 4; scale_cache[0] = scale_vec[0]; scale_cache[1] = scale_vec[1]; scale_cache[2] = scale_vec[2]; scale_cache[3] = scale_vec[3]; - half2 const* zp_vec = reinterpret_cast( - zp_expert_base_ + metadata_offset); + half2 const* zp_vec = + reinterpret_cast(zp_expert_base_ + metadata_offset); half2* zp_cache = cached_zp_ + c * 4; zp_cache[0] = zp_vec[0]; zp_cache[1] = zp_vec[1]; @@ -181,8 +179,8 @@ class Sm70MoeU8ZpIteratorB { CUTLASS_DEVICE void cache_metadata_vector_words(int c, int group, int cache_n) const { int const metadata_offset = metadata_group_offset(group) + cache_n; - uint4 const scale_words = *reinterpret_cast( - scales_expert_base_ + metadata_offset); + uint4 const scale_words = + *reinterpret_cast(scales_expert_base_ + metadata_offset); half2 const* scale_vec = reinterpret_cast(&scale_words); half2* scale_cache = cached_scales_ + c * 4; scale_cache[0] = scale_vec[0]; @@ -190,8 +188,8 @@ class Sm70MoeU8ZpIteratorB { scale_cache[2] = scale_vec[2]; scale_cache[3] = scale_vec[3]; - uint4 const zp_words = *reinterpret_cast( - zp_expert_base_ + metadata_offset); + uint4 const zp_words = + *reinterpret_cast(zp_expert_base_ + metadata_offset); half2 const* zp_vec = reinterpret_cast(&zp_words); half2* zp_cache = cached_zp_ + c * 4; zp_cache[0] = zp_vec[0]; @@ -204,9 +202,8 @@ class Sm70MoeU8ZpIteratorB { void cache_current_group_metadata(int group) const { CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; if constexpr (kUseMetadataVectorWords) { cache_metadata_vector_words(c, group, cache_n); } else { @@ -269,8 +266,9 @@ class Sm70MoeU8ZpIteratorB { frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE U8 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE U8 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_offset_); uint32_t const qword1 = load_qword_vector<1>( @@ -280,12 +278,12 @@ class Sm70MoeU8ZpIteratorB { half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data()); - marlin::dequant( - static_cast(qword0), deq); + marlin::dequant(static_cast(qword0), + deq); frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); - marlin::dequant( - static_cast(qword1), deq); + marlin::dequant(static_cast(qword1), + deq); frag_vec[2] = __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } @@ -293,9 +291,8 @@ class Sm70MoeU8ZpIteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; if constexpr (kGroupSize != -1) { cache_current_group_metadata(scale_group(logical_k_s)); } @@ -303,10 +300,9 @@ class Sm70MoeU8ZpIteratorB { expert_qweight_base_offset() + qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - uint4 const qwords0 = - load_qword_vector<4>(qweight_ + qweight_base_s); - uint4 const qwords1 = load_qword_vector<4>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint4 const qwords0 = load_qword_vector<4>(qweight_ + qweight_base_s); + uint4 const qwords1 = load_qword_vector<4>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -321,22 +317,17 @@ class Sm70MoeU8ZpIteratorB { half2* frag_vec = reinterpret_cast(frag.data() + frag_base); marlin::dequant( static_cast(qword0), deq); - frag_vec[0] = - __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); - frag_vec[1] = - __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); + frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); + frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( static_cast(qword1), deq); - frag_vec[2] = - __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); - frag_vec[3] = - __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); + frag_vec[2] = __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); + frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - uint2 const qwords0 = - load_qword_vector<2>(qweight_ + qweight_base_s); - uint2 const qwords1 = load_qword_vector<2>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint2 const qwords0 = load_qword_vector<2>(qweight_ + qweight_base_s); + uint2 const qwords1 = load_qword_vector<2>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -351,20 +342,17 @@ class Sm70MoeU8ZpIteratorB { half2* frag_vec = reinterpret_cast(frag.data() + frag_base); marlin::dequant( static_cast(qword0), deq); - frag_vec[0] = - __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); - frag_vec[1] = - __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); + frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); + frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( static_cast(qword1), deq); - frag_vec[2] = - __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); - frag_vec[3] = - __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); + frag_vec[2] = __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); + frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE U8 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE U8 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_s); uint32_t const qword1 = load_qword_vector<1>( @@ -395,8 +383,7 @@ class Sm70MoeU8ZpIteratorB { return; } - if constexpr (kGroupSize != -1 && - ThreadMap::Iterations::kStrided == 1) { + if constexpr (kGroupSize != -1 && ThreadMap::Iterations::kStrided == 1) { int const first_logical_k = k_offset_ + thread_offset_.strided(); cache_current_group_metadata(scale_group(first_logical_k)); } @@ -415,18 +402,17 @@ struct Sm70MoeU8ZpGemmSpec { using IteratorA = Sm70MoeGatherIteratorA; template - using IteratorB = - Sm70MoeU8ZpIteratorB; + using IteratorB = Sm70MoeU8ZpIteratorB; }; -template using Sm70MoeU8ZpGemmTraits = - Sm70MarlinMoeGemmTraits< - Sm70MoeU8ZpGemmSpec, CtaM, CtaN, CtaK, - Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; + Sm70MarlinMoeGemmTraits, CtaM, + CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>; template struct Sm70MoeU8Launcher { @@ -448,16 +434,16 @@ struct Sm70MoeU8Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - using Traits = Sm70MoeU8ZpGemmTraits; + using Traits = + Sm70MoeU8ZpGemmTraits; return launch_sm70_marlin_moe_gemm( a, c, b_q_weight, b_scales, b_zeros, global_scale, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, - top_k, mul_topk_weights, size_m, size_n, size_k, requested_split_k); + expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, + mul_topk_weights, size_m, size_n, size_k, requested_split_k); } }; @@ -475,9 +461,10 @@ torch::Tensor sm70_marlin_u8_gemm( auto const params = sm70_marlin_moe_auto_stage_params( "uint8", group_size, moe_block_size, top_k, size_m, size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE uint8", geometry); - validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE uint8", geometry, - size_n); + validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE uint8", + geometry); + validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE uint8", + geometry, size_n); TORCH_CHECK(size_k % geometry.cta_k == 0, "SM70 Marlin MoE uint8 requires K divisible by CTA_K=", geometry.cta_k, ". Got K=", size_k, "."); @@ -485,19 +472,45 @@ torch::Tensor sm70_marlin_u8_gemm( auto empty_float = torch::empty( {0}, torch::TensorOptions().dtype(at::kFloat).device(a.device())); if (params.use_metadata_vector_words) { - Sm70MoeU8Launcher const launcher{ - a, c, b_q_weight, b_scales, b_zeros, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; - return dispatch_sm70_marlin_moe_geometry_with_group32(launcher, geometry, params.packed_macro_n, group_size, - "uint8"); + Sm70MoeU8Launcher const launcher{a, + c, + b_q_weight, + b_scales, + b_zeros, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; + return dispatch_sm70_marlin_moe_geometry_with_group32( + launcher, geometry, params.packed_macro_n, group_size, "uint8"); } - Sm70MoeU8Launcher const launcher{ - a, c, b_q_weight, b_scales, b_zeros, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; - return dispatch_sm70_marlin_moe_geometry_with_group32(launcher, geometry, params.packed_macro_n, group_size, - "uint8"); + Sm70MoeU8Launcher const launcher{a, + c, + b_q_weight, + b_scales, + b_zeros, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; + return dispatch_sm70_marlin_moe_geometry_with_group32( + launcher, geometry, params.packed_macro_n, group_size, "uint8"); } } // namespace marlin_moe_wna16 diff --git a/csrc/moe/marlin_moe_wna16/sm70_marlin_u8b128_gemm.cu b/csrc/moe/marlin_moe_wna16/sm70_marlin_u8b128_gemm.cu index 511e4d4181..1eabec24dc 100644 --- a/csrc/moe/marlin_moe_wna16/sm70_marlin_u8b128_gemm.cu +++ b/csrc/moe/marlin_moe_wna16/sm70_marlin_u8b128_gemm.cu @@ -35,22 +35,25 @@ class Sm70MoeU8B128IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; - static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, - "SM70 Marlin MoE U8B128 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, - "SM70 Marlin MoE U8B128 IteratorB expects one contiguous iteration per " - "64-column quant tile."); + using Fragment = cutlass::Array; + static_assert( + Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, + "SM70 Marlin MoE U8B128 IteratorB expects CTA_N in {64, 128, 256}."); + static_assert( + ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, + "SM70 Marlin MoE U8B128 IteratorB expects one contiguous iteration per " + "64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, "SM70 Marlin MoE U8B128 IteratorB expects 64-column deltas."); - static_assert(ThreadMap::kElementsPerAccess == kU8B128ValuesPerAccess, - "SM70 Marlin MoE U8B128 IteratorB expects two packed U8B128 words " - "per access."); - static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin MoE U8-family IteratorB expects one or more strided " - "iterations."); + static_assert( + ThreadMap::kElementsPerAccess == kU8B128ValuesPerAccess, + "SM70 Marlin MoE U8B128 IteratorB expects two packed U8B128 words " + "per access."); + static_assert( + ThreadMap::Iterations::kStrided >= 1, + "SM70 Marlin MoE U8-family IteratorB expects one or more strided " + "iterations."); static constexpr int kQweightWordStrideWords = kPackedMacroN / kQuantTileN; struct Params { @@ -86,9 +89,8 @@ class Sm70MoeU8B128IteratorB { half const* scales, half const*, int thread_id, int expert, int k_offset, int n_offset) : qweight_(qweight), - scales_expert_base_(scales + - (expert >= 0 ? expert : 0) * - params.num_groups * params.size_n), + scales_expert_base_(scales + (expert >= 0 ? expert : 0) * + params.num_groups * params.size_n), params_(params), thread_offset_(ThreadMap::initial_offset(thread_id)), expert_(expert), @@ -130,19 +132,17 @@ class Sm70MoeU8B128IteratorB { } CUTLASS_DEVICE - int metadata_group_offset(int group) const { - return group * params_.size_n; - } + int metadata_group_offset(int group) const { return group * params_.size_n; } CUTLASS_DEVICE int scale_group(int logical_k) const { if constexpr (kGroupSize == -1) { return 0; } else { - static_assert(kGroupSize == 32 || kGroupSize == 64 || - kGroupSize == 128, - "SM70 Marlin MoE U8B128 supports only group sizes -1, 32, 64, " - "and 128."); + static_assert( + kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, + "SM70 Marlin MoE U8B128 supports only group sizes -1, 32, 64, " + "and 128."); return logical_k / kGroupSize; } } @@ -157,8 +157,8 @@ class Sm70MoeU8B128IteratorB { CUTLASS_DEVICE void cache_metadata_lane_vectors(int c, int group, int cache_n) const { int const metadata_offset = metadata_group_offset(group) + cache_n; - half2 const* scale_vec = reinterpret_cast( - scales_expert_base_ + metadata_offset); + half2 const* scale_vec = + reinterpret_cast(scales_expert_base_ + metadata_offset); half2* scale_cache = cached_scales_ + c * 4; scale_cache[0] = scale_vec[0]; scale_cache[1] = scale_vec[1]; @@ -169,8 +169,8 @@ class Sm70MoeU8B128IteratorB { CUTLASS_DEVICE void cache_metadata_vector_words(int c, int group, int cache_n) const { int const metadata_offset = metadata_group_offset(group) + cache_n; - uint4 const scale_words = *reinterpret_cast( - scales_expert_base_ + metadata_offset); + uint4 const scale_words = + *reinterpret_cast(scales_expert_base_ + metadata_offset); half2 const* scale_vec = reinterpret_cast(&scale_words); half2* scale_cache = cached_scales_ + c * 4; scale_cache[0] = scale_vec[0]; @@ -183,9 +183,8 @@ class Sm70MoeU8B128IteratorB { void cache_current_group_metadata(int group) const { CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; if constexpr (kUseMetadataVectorWords) { cache_metadata_vector_words(c, group, cache_n); } else { @@ -246,8 +245,9 @@ class Sm70MoeU8B128IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE U8B128 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE U8B128 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_offset_); uint32_t const qword1 = load_qword_vector<1>( @@ -269,9 +269,8 @@ class Sm70MoeU8B128IteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; if constexpr (kGroupSize != -1) { cache_current_group_metadata(scale_group(logical_k_s)); } @@ -279,10 +278,9 @@ class Sm70MoeU8B128IteratorB { expert_qweight_base_offset() + qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - uint4 const qwords0 = - load_qword_vector<4>(qweight_ + qweight_base_s); - uint4 const qwords1 = load_qword_vector<4>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint4 const qwords0 = load_qword_vector<4>(qweight_ + qweight_base_s); + uint4 const qwords1 = load_qword_vector<4>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -296,22 +294,17 @@ class Sm70MoeU8B128IteratorB { half2* frag_vec = reinterpret_cast(frag.data() + frag_base); marlin::dequant( static_cast(qword0), deq); - frag_vec[0] = - __hmul2(deq[0], scale_vec[0]); - frag_vec[1] = - __hmul2(deq[1], scale_vec[1]); + frag_vec[0] = __hmul2(deq[0], scale_vec[0]); + frag_vec[1] = __hmul2(deq[1], scale_vec[1]); marlin::dequant( static_cast(qword1), deq); - frag_vec[2] = - __hmul2(deq[0], scale_vec[2]); - frag_vec[3] = - __hmul2(deq[1], scale_vec[3]); + frag_vec[2] = __hmul2(deq[0], scale_vec[2]); + frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - uint2 const qwords0 = - load_qword_vector<2>(qweight_ + qweight_base_s); - uint2 const qwords1 = load_qword_vector<2>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint2 const qwords0 = load_qword_vector<2>(qweight_ + qweight_base_s); + uint2 const qwords1 = load_qword_vector<2>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -325,20 +318,17 @@ class Sm70MoeU8B128IteratorB { half2* frag_vec = reinterpret_cast(frag.data() + frag_base); marlin::dequant( static_cast(qword0), deq); - frag_vec[0] = - __hmul2(deq[0], scale_vec[0]); - frag_vec[1] = - __hmul2(deq[1], scale_vec[1]); + frag_vec[0] = __hmul2(deq[0], scale_vec[0]); + frag_vec[1] = __hmul2(deq[1], scale_vec[1]); marlin::dequant( static_cast(qword1), deq); - frag_vec[2] = - __hmul2(deq[0], scale_vec[2]); - frag_vec[3] = - __hmul2(deq[1], scale_vec[3]); + frag_vec[2] = __hmul2(deq[0], scale_vec[2]); + frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MoE U8B128 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MoE U8B128 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_s); uint32_t const qword1 = load_qword_vector<1>( @@ -368,8 +358,7 @@ class Sm70MoeU8B128IteratorB { return; } - if constexpr (kGroupSize != -1 && - ThreadMap::Iterations::kStrided == 1) { + if constexpr (kGroupSize != -1 && ThreadMap::Iterations::kStrided == 1) { int const first_logical_k = k_offset_ + thread_offset_.strided(); cache_current_group_metadata(scale_group(first_logical_k)); } @@ -393,13 +382,13 @@ struct Sm70MoeU8B128GemmSpec { UseMetadataVectorWords>; }; -template using Sm70MoeU8B128GemmTraits = - Sm70MarlinMoeGemmTraits< - Sm70MoeU8B128GemmSpec, CtaM, CtaN, CtaK, - Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; + Sm70MarlinMoeGemmTraits, CtaM, + CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>; template struct Sm70MoeU8B128Launcher { @@ -421,16 +410,17 @@ struct Sm70MoeU8B128Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - using Traits = Sm70MoeU8B128GemmTraits; + using Traits = + Sm70MoeU8B128GemmTraits; return launch_sm70_marlin_moe_gemm( a, c, b_q_weight, b_scales, b_zeros, global_scale, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, - top_k, mul_topk_weights, size_m, size_n, size_k, requested_split_k); + expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, + mul_topk_weights, size_m, size_n, size_k, requested_split_k); } }; @@ -448,9 +438,10 @@ torch::Tensor sm70_marlin_u8b128_gemm( auto const params = sm70_marlin_moe_auto_stage_params( "uint8b128", group_size, moe_block_size, top_k, size_m, size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_moe_stage_cta_geometry_supported("SM70 Marlin MoE uint8b128", geometry); - validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE uint8b128", geometry, - size_n); + validate_sm70_marlin_moe_stage_cta_geometry_supported( + "SM70 Marlin MoE uint8b128", geometry); + validate_sm70_marlin_moe_stage_cta_n_alignment("SM70 Marlin MoE uint8b128", + geometry, size_n); TORCH_CHECK(size_k % geometry.cta_k == 0, "SM70 Marlin MoE uint8b128 requires K divisible by CTA_K=", geometry.cta_k, ". Got K=", size_k, "."); @@ -459,17 +450,43 @@ torch::Tensor sm70_marlin_u8b128_gemm( auto empty_float = torch::empty( {0}, torch::TensorOptions().dtype(at::kFloat).device(a.device())); if (params.use_metadata_vector_words) { - Sm70MoeU8B128Launcher const launcher{ - a, c, b_q_weight, b_scales, empty_half, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; + Sm70MoeU8B128Launcher const launcher{a, + c, + b_q_weight, + b_scales, + empty_half, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; return dispatch_sm70_marlin_moe_geometry_with_group32( launcher, geometry, params.packed_macro_n, group_size, "uint8b128"); } - Sm70MoeU8B128Launcher const launcher{ - a, c, b_q_weight, b_scales, empty_half, empty_float, sorted_token_ids, - expert_ids, num_tokens_past_padded, topk_weights, moe_block_size, top_k, - mul_topk_weights, size_m, size_n, size_k, params.requested_split_k}; + Sm70MoeU8B128Launcher const launcher{a, + c, + b_q_weight, + b_scales, + empty_half, + empty_float, + sorted_token_ids, + expert_ids, + num_tokens_past_padded, + topk_weights, + moe_block_size, + top_k, + mul_topk_weights, + size_m, + size_n, + size_k, + params.requested_split_k}; return dispatch_sm70_marlin_moe_geometry_with_group32( launcher, geometry, params.packed_macro_n, group_size, "uint8b128"); } diff --git a/csrc/moe/moe_permute_unpermute_op.cu b/csrc/moe/moe_permute_unpermute_op.cu index 9cd84897ff..334629cc98 100644 --- a/csrc/moe/moe_permute_unpermute_op.cu +++ b/csrc/moe/moe_permute_unpermute_op.cu @@ -28,9 +28,7 @@ torch::Tensor maybe_allocate_tensor( return torch::empty(expected_sizes, torch::dtype(dtype).device(device)); } -int64_t pad_to_multiple_of_16(int64_t value) { - return (value + 15) / 16 * 16; -} +int64_t pad_to_multiple_of_16(int64_t value) { return (value + 15) / 16 * 16; } } // namespace @@ -72,16 +70,15 @@ void moe_permute_impl( auto expanded_rows = n_token * topk; auto stream = at::cuda::getCurrentCUDAStream().stream(); - if (canUseSingleTokenMoePermuteFastPath( - n_token, topk, expert_map.has_value(), n_expert, n_local_expert)) { + if (canUseSingleTokenMoePermuteFastPath(n_token, topk, expert_map.has_value(), + n_expert, n_local_expert)) { if (input.scalar_type() == at::ScalarType::Half) { singleTokenMoePermuteLauncher( get_ptr(input), get_ptr(topk_ids), get_ptr(permuted_input), get_ptr(expert_first_token_offset), get_ptr(inv_permuted_idx), get_ptr(permuted_idx), - static_cast(n_expert), static_cast(topk), n_hidden, - stream); + static_cast(n_expert), static_cast(topk), n_hidden, stream); return; } if (input.scalar_type() == at::ScalarType::BFloat16) { @@ -90,8 +87,7 @@ void moe_permute_impl( get_ptr<__nv_bfloat16>(permuted_input), get_ptr(expert_first_token_offset), get_ptr(inv_permuted_idx), get_ptr(permuted_idx), - static_cast(n_expert), static_cast(topk), n_hidden, - stream); + static_cast(n_expert), static_cast(topk), n_hidden, stream); return; } if (input.scalar_type() == at::ScalarType::Float) { @@ -100,8 +96,7 @@ void moe_permute_impl( get_ptr(permuted_input), get_ptr(expert_first_token_offset), get_ptr(inv_permuted_idx), get_ptr(permuted_idx), - static_cast(n_expert), static_cast(topk), n_hidden, - stream); + static_cast(n_expert), static_cast(topk), n_hidden, stream); return; } } diff --git a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu index bf86c12abb..eaf633af6b 100644 --- a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu @@ -93,10 +93,12 @@ bool singleTokenUnpermuteFastPathEnabled() { } template -__global__ void singleTokenMoePermuteKernel( - T const* input, int const* topk_ids, T* permuted_output, - int64_t* expert_first_token_offset, int* inv_permuted_idx, - int* permuted_idx, int num_experts, int topk, int64_t cols) { +__global__ void singleTokenMoePermuteKernel(T const* input, int const* topk_ids, + T* permuted_output, + int64_t* expert_first_token_offset, + int* inv_permuted_idx, + int* permuted_idx, int num_experts, + int topk, int64_t cols) { __shared__ int sorted_ids[kSingleTokenFastPathMaxTopK]; __shared__ int sorted_src[kSingleTokenFastPathMaxTopK]; @@ -177,8 +179,8 @@ __global__ void singleTokenMoeUnpermuteKernel( auto const* expanded_rows_v = reinterpret_cast(expanded_permuted_rows); - auto* reduced_row_ptr_v = reinterpret_cast( - reduced_unpermuted_output); + auto* reduced_row_ptr_v = + reinterpret_cast(reduced_unpermuted_output); int64_t const num_elems_in_col = cols / FINALIZE_ELEM_PER_THREAD; for (int64_t elem_index = threadIdx.x; elem_index < num_elems_in_col; @@ -186,7 +188,7 @@ __global__ void singleTokenMoeUnpermuteKernel( ComputeElem thread_output; thread_output.fill(0); -#pragma unroll + #pragma unroll for (int k_idx = 0; k_idx < kSingleTokenFastPathMaxTopK; ++k_idx) { if (k_idx >= topk) { break; @@ -222,11 +224,12 @@ bool canUseSingleTokenMoePermuteFastPath(int64_t n_token, int64_t topk, } template -void singleTokenMoePermuteLauncher( - T const* input, int const* topk_ids, T* permuted_output, - int64_t* expert_first_token_offset, int* inv_permuted_idx, - int* permuted_idx, int num_experts, int topk, int64_t cols, - cudaStream_t stream) { +void singleTokenMoePermuteLauncher(T const* input, int const* topk_ids, + T* permuted_output, + int64_t* expert_first_token_offset, + int* inv_permuted_idx, int* permuted_idx, + int num_experts, int topk, int64_t cols, + cudaStream_t stream) { constexpr int threads = 256; singleTokenMoePermuteKernel<<<1, threads, 0, stream>>>( input, topk_ids, permuted_output, expert_first_token_offset, diff --git a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h index d87ed201de..fcce90dad4 100644 --- a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h @@ -80,11 +80,12 @@ bool canUseSingleTokenMoePermuteFastPath(int64_t n_token, int64_t topk, int64_t n_local_expert); template -void singleTokenMoePermuteLauncher( - T const* input, int const* topk_ids, T* permuted_output, - int64_t* expert_first_token_offset, int* inv_permuted_idx, - int* permuted_idx, int num_experts, int topk, int64_t cols, - cudaStream_t stream); +void singleTokenMoePermuteLauncher(T const* input, int const* topk_ids, + T* permuted_output, + int64_t* expert_first_token_offset, + int* inv_permuted_idx, int* permuted_idx, + int num_experts, int topk, int64_t cols, + cudaStream_t stream); bool canUseSingleTokenMoeUnpermuteFastPath(int64_t n_token, int64_t topk); diff --git a/csrc/ops.h b/csrc/ops.h index af458e78dc..1588001387 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -193,6 +193,13 @@ void fp8_qpn8_gated_pair_sm70_out(torch::Tensor out, torch::Tensor input, int64_t accumulator_chains, bool fast_decoder, bool prefetch_codes); +void fp8_qpn8_hc_dispatch_sm70_out( + torch::Tensor block_out, torch::Tensor injection_out, + torch::Tensor down_staging, torch::Tensor lora_staging, + torch::Tensor gate_staging, torch::Tensor partials, + int64_t dense_weight_ptr, torch::Tensor xn, torch::Tensor down_codes, + torch::Tensor down_scales, torch::Tensor up_codes, torch::Tensor up_scales); + std::vector nvfp4_qpn4_prepare_sm70(torch::Tensor qweight, torch::Tensor scales); @@ -220,6 +227,12 @@ void fp8_gemm_sm70_prefill_prescaled_out(torch::Tensor out, int64_t group_size, int64_t k_ld, int64_t q_ld); +void fp8_gemm_sm70_prescaled_m1_out(torch::Tensor out, torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _prescaled_factors, + int64_t group_size, int64_t k_ld, + int64_t q_ld); + std::vector nvfp4_qpn2_prepare_sm70(torch::Tensor weight_packed, torch::Tensor weight_scale); @@ -308,6 +321,18 @@ void sm70_f16_rerank_topk_out(torch::Tensor values_out, torch::Tensor ids_out, torch::Tensor logits, torch::Tensor candidate_ids, int64_t vocab_start_index); +void sm70_glm_mhc_pre_norm_out( + torch::Tensor gemm_mul, torch::Tensor gemm_sqrsum, torch::Tensor hc_scale, + torch::Tensor hc_base, torch::Tensor residual, torch::Tensor post_mix, + torch::Tensor comb_mix, torch::Tensor layer_input, + torch::Tensor norm_weight, double rms_eps, double hc_pre_eps, + double hc_sinkhorn_eps, double hc_post_mult, int64_t sinkhorn_repeat, + double norm_eps); + +void sm70_glm_kda_fg_b_out(torch::Tensor f_out, torch::Tensor g_out, + torch::Tensor f_input, torch::Tensor g_input, + torch::Tensor f_weight, torch::Tensor g_weight); + void sm70_f16_lm_head_top1_out(torch::Tensor values_out, torch::Tensor indices_out, torch::Tensor _in_feats, torch::Tensor _kernel, @@ -479,6 +504,15 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, int64_t num_experts, int64_t k, int64_t n, int64_t group_size); +void mxfp4_moe_qpn_m1_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, torch::Tensor scales, + torch::Tensor expert_ids, bool broadcast_input); + +void nvfp4_moe_qpn_m1_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, torch::Tensor scales, + torch::Tensor expert_ids, bool broadcast_input, + int64_t split_k); + void nvfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, torch::Tensor expert_offsets, torch::Tensor dense_expert_ids, diff --git a/csrc/quantization/marlin/sm70_awq_marlin_repack.cu b/csrc/quantization/marlin/sm70_awq_marlin_repack.cu index ff37c46ff5..b664b2b29f 100644 --- a/csrc/quantization/marlin/sm70_awq_marlin_repack.cu +++ b/csrc/quantization/marlin/sm70_awq_marlin_repack.cu @@ -92,8 +92,7 @@ __global__ void awq_marlin_repack_kernel( constexpr int sh_stride = tile_n_ints; int4* sh_stage_ptr = sh + stage_size * pipe; uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); - uint32_t packed_src = - sh_stage_int_ptr[local_k * sh_stride + local_n_vec]; + uint32_t packed_src = sh_stage_int_ptr[local_k * sh_stride + local_n_vec]; constexpr int undo_pack[8] = {0, 4, 1, 5, 2, 6, 3, 7}; constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; @@ -114,9 +113,8 @@ __global__ void awq_marlin_repack_kernel( int const cta_first_n_tile = cta_n_tile * cta_n_tiles; int const subtile = n_tile_id - cta_first_n_tile; int const local_word = local_k * 8 + local_n_vec; - int const cta_n_offset = - cta_n_tile * cta_n_tiles * tile_size + - local_word * cta_n_tiles + subtile; + int const cta_n_offset = cta_n_tile * cta_n_tiles * tile_size + + local_word * cta_n_tiles + subtile; out_ptr[k_tile_id * n_tiles * tile_size + cta_n_offset] = res; return; } @@ -155,9 +153,8 @@ __global__ void awq_marlin_repack_kernel( int const cta_first_n_tile = cta_n_tile * cta_n_tiles; int const subtile = n_tile_id - cta_first_n_tile; int const local_word = local_k * 16 + local_n_word; - int const cta_n_offset = - cta_n_tile * cta_n_tiles * tile_size + - local_word * cta_n_tiles + subtile; + int const cta_n_offset = cta_n_tile * cta_n_tiles * tile_size + + local_word * cta_n_tiles + subtile; out_ptr[k_tile_id * n_tiles * tile_size + cta_n_offset] = res; return; } @@ -304,16 +301,15 @@ __global__ void awq_marlin_repack_kernel( b_q_weight_ptr, out_ptr, size_k, size_n); \ } else -#define CALL_FOR_CTA(CTA_N) \ - do { \ - CALL_IF(4, false, CTA_N) \ - CALL_IF(8, false, CTA_N) \ - CALL_IF(4, true, CTA_N) \ - CALL_IF(8, true, CTA_N) \ - { \ - TORCH_CHECK(false, "Unsupported repack config: num_bits = ", \ - num_bits, ", is_a_8bit = ", is_a_8bit); \ - } \ +#define CALL_FOR_CTA(CTA_N) \ + do { \ + CALL_IF(4, false, CTA_N) \ + CALL_IF(8, false, CTA_N) \ + CALL_IF(4, true, CTA_N) \ + CALL_IF(8, true, CTA_N) { \ + TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, \ + ", is_a_8bit = ", is_a_8bit); \ + } \ } while (false) torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k, diff --git a/csrc/quantization/marlin/sm70_gptq_marlin_repack.cu b/csrc/quantization/marlin/sm70_gptq_marlin_repack.cu index db89b2d5af..089ea3220d 100644 --- a/csrc/quantization/marlin/sm70_gptq_marlin_repack.cu +++ b/csrc/quantization/marlin/sm70_gptq_marlin_repack.cu @@ -164,9 +164,8 @@ __global__ void gptq_marlin_repack_kernel( int const cta_first_n_tile = cta_n_tile * cta_n_tiles; int const subtile = n_tile_id - cta_first_n_tile; int const local_word = local_k * 8 + local_n_vec; - int const cta_n_offset = - cta_n_tile * cta_n_tiles * tile_size + - local_word * cta_n_tiles + subtile; + int const cta_n_offset = cta_n_tile * cta_n_tiles * tile_size + + local_word * cta_n_tiles + subtile; out_ptr[k_tile_id * n_tiles * tile_size + cta_n_offset] = res; return; } @@ -215,9 +214,8 @@ __global__ void gptq_marlin_repack_kernel( int const cta_first_n_tile = cta_n_tile * cta_n_tiles; int const subtile = n_tile_id - cta_first_n_tile; int const local_word = local_k * 16 + local_n_word; - int const cta_n_offset = - cta_n_tile * cta_n_tiles * tile_size + - local_word * cta_n_tiles + subtile; + int const cta_n_offset = cta_n_tile * cta_n_tiles * tile_size + + local_word * cta_n_tiles + subtile; out_ptr[k_tile_id * n_tiles * tile_size + cta_n_offset] = res; return; } @@ -381,19 +379,17 @@ __global__ void gptq_marlin_repack_kernel( b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ } else -#define CALL_FOR_CTA(CTA_N) \ - do { \ - CALL_IF(4, false, false, CTA_N) \ - CALL_IF(4, true, false, CTA_N) \ - CALL_IF(8, false, false, CTA_N) \ - CALL_IF(8, true, false, CTA_N) \ - CALL_IF(4, false, true, CTA_N) \ - CALL_IF(8, false, true, CTA_N) \ - { \ - TORCH_CHECK(false, "Unsupported repack config: num_bits = ", \ - num_bits, ", has_perm = ", has_perm, \ - ", is_a_8bit = ", is_a_8bit); \ - } \ +#define CALL_FOR_CTA(CTA_N) \ + do { \ + CALL_IF(4, false, false, CTA_N) \ + CALL_IF(4, true, false, CTA_N) \ + CALL_IF(8, false, false, CTA_N) \ + CALL_IF(8, true, false, CTA_N) \ + CALL_IF(4, false, true, CTA_N) \ + CALL_IF(8, false, true, CTA_N) { \ + TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, \ + ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit); \ + } \ } while (false) torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, diff --git a/csrc/quantization/marlin/sm70_marlin_common.cuh b/csrc/quantization/marlin/sm70_marlin_common.cuh index 0e8ecbf4a4..3ae26c8fcf 100644 --- a/csrc/quantization/marlin/sm70_marlin_common.cuh +++ b/csrc/quantization/marlin/sm70_marlin_common.cuh @@ -28,31 +28,29 @@ inline size_t configure_sm70_dynamic_smem(Kernel kernel) { return smem_bytes; } -constexpr int sm70_const_min(int lhs, int rhs) { - return lhs < rhs ? lhs : rhs; -} +constexpr int sm70_const_min(int lhs, int rhs) { return lhs < rhs ? lhs : rhs; } -constexpr bool sm70_marlin_basic_geometry_values_supported( - int cta_m, int cta_n, int cta_k, int warps, int warp_m, int warp_n, - int warp_k) { +constexpr bool sm70_marlin_basic_geometry_values_supported(int cta_m, int cta_n, + int cta_k, int warps, + int warp_m, + int warp_n, + int warp_k) { return (cta_m == 32 || cta_m == 64 || cta_m == 128 || cta_m == 256) && (cta_n == 64 || cta_n == 128 || cta_n == 256) && (cta_k == 16 || cta_k == 32 || cta_k == 64 || cta_k == 128) && - (warps == 4 || warps == 8) && - (warp_m == 32 || warp_m == 64) && - (warp_n == 32 || warp_n == 64) && - (warp_k == 16 || warp_k == 32); + (warps == 4 || warps == 8) && (warp_m == 32 || warp_m == 64) && + (warp_n == 32 || warp_n == 64) && (warp_k == 16 || warp_k == 32); } -constexpr bool sm70_marlin_warp_decomposition_supported( - int cta_m, int cta_n, int cta_k, int warps, int warp_m, int warp_n, - int warp_k) { - if (!sm70_marlin_basic_geometry_values_supported( - cta_m, cta_n, cta_k, warps, warp_m, warp_n, warp_k)) { +constexpr bool sm70_marlin_warp_decomposition_supported(int cta_m, int cta_n, + int cta_k, int warps, + int warp_m, int warp_n, + int warp_k) { + if (!sm70_marlin_basic_geometry_values_supported(cta_m, cta_n, cta_k, warps, + warp_m, warp_n, warp_k)) { return false; } - if (cta_m % warp_m != 0 || cta_n % warp_n != 0 || - cta_k % warp_k != 0) { + if (cta_m % warp_m != 0 || cta_n % warp_n != 0 || cta_k % warp_k != 0) { return false; } return (cta_m / warp_m) * (cta_n / warp_n) * (cta_k / warp_k) == warps; @@ -65,8 +63,7 @@ constexpr bool sm70_marlin_pitchlinear_threadmap_supported( if (shape_contiguous % elements_per_access != 0) { return false; } - int const shape_access_contiguous = - shape_contiguous / elements_per_access; + int const shape_access_contiguous = shape_contiguous / elements_per_access; int const shape_access_strided = shape_strided; if (shape_access_contiguous % warp_thread_contiguous != 0 || shape_access_strided % warp_thread_strided != 0) { @@ -74,8 +71,7 @@ constexpr bool sm70_marlin_pitchlinear_threadmap_supported( } int const warp_access_contiguous = shape_access_contiguous / warp_thread_contiguous; - int const warp_access_strided = - shape_access_strided / warp_thread_strided; + int const warp_access_strided = shape_access_strided / warp_thread_strided; int const warp_count = threads / 32; if (warp_access_strided <= 0 || warp_count <= 0) { return false; @@ -90,43 +86,40 @@ constexpr bool sm70_marlin_pitchlinear_threadmap_supported( if (warps_contiguous <= 0) { return false; } - int const iterations_contiguous = - warp_access_contiguous / warps_contiguous; + int const iterations_contiguous = warp_access_contiguous / warps_contiguous; int const iterations_strided = warp_access_strided / warps_strided; return iterations_contiguous > 0 && iterations_strided > 0; } -constexpr bool sm70_marlin_b_threadmap_matches_quant_iterator( - int cta_n, int cta_k, int threads) { +constexpr bool sm70_marlin_b_threadmap_matches_quant_iterator(int cta_n, + int cta_k, + int threads) { constexpr int kElementsPerAccess = 8; constexpr int kWarpThreadContiguous = 8; constexpr int kWarpThreadStrided = 4; if (!sm70_marlin_pitchlinear_threadmap_supported( - cta_n, cta_k, threads, kWarpThreadContiguous, - kWarpThreadStrided, kElementsPerAccess)) { + cta_n, cta_k, threads, kWarpThreadContiguous, kWarpThreadStrided, + kElementsPerAccess)) { return false; } int const shape_access_contiguous = cta_n / kElementsPerAccess; int const shape_access_strided = cta_k; int const warp_access_contiguous = shape_access_contiguous / kWarpThreadContiguous; - int const warp_access_strided = - shape_access_strided / kWarpThreadStrided; + int const warp_access_strided = shape_access_strided / kWarpThreadStrided; int const warp_count = threads / 32; int const warps_strided = warp_access_strided >= warp_count ? warp_count : warp_access_strided; int const warps_contiguous = warp_count > warp_access_strided ? warp_count / warps_strided : 1; - int const iterations_contiguous = - warp_access_contiguous / warps_contiguous; - int const delta_contiguous = - kWarpThreadContiguous * kElementsPerAccess; + int const iterations_contiguous = warp_access_contiguous / warps_contiguous; + int const delta_contiguous = kWarpThreadContiguous * kElementsPerAccess; return iterations_contiguous == cta_n / kQuantTileN && delta_contiguous == kQuantTileN; } -constexpr bool sm70_marlin_volta_epilogue_supported( - int cta_m, int cta_n, int warps, int warp_m) { +constexpr bool sm70_marlin_volta_epilogue_supported(int cta_m, int cta_n, + int warps, int warp_m) { constexpr int kShapeRow = 4; constexpr int kShapeGroup = 4; constexpr int kElementsPerAccess = 8; @@ -183,8 +176,8 @@ constexpr bool sm70_marlin_volta_epilogue_supported( constexpr bool sm70_marlin_cta_geometry_is_supported_constexpr( int cta_m, int cta_n, int cta_k, int warps, int warp_m, int warp_n, int warp_k) { - if (!sm70_marlin_warp_decomposition_supported( - cta_m, cta_n, cta_k, warps, warp_m, warp_n, warp_k)) { + if (!sm70_marlin_warp_decomposition_supported(cta_m, cta_n, cta_k, warps, + warp_m, warp_n, warp_k)) { return false; } int const threads = warps * 32; @@ -192,12 +185,11 @@ constexpr bool sm70_marlin_cta_geometry_is_supported_constexpr( // DefaultMmaCore. Its A/B thread maps are fixed to 128-bit half loads. // CTA_K=16 is a mathematical GEMM shape but does not provide enough // contiguous A accesses for that thread map. - if (!sm70_marlin_pitchlinear_threadmap_supported( - cta_k, cta_m, threads, 4, 8, 8)) { + if (!sm70_marlin_pitchlinear_threadmap_supported(cta_k, cta_m, threads, 4, 8, + 8)) { return false; } - if (!sm70_marlin_b_threadmap_matches_quant_iterator(cta_n, cta_k, - threads)) { + if (!sm70_marlin_b_threadmap_matches_quant_iterator(cta_n, cta_k, threads)) { return false; } return sm70_marlin_volta_epilogue_supported(cta_m, cta_n, warps, warp_m); @@ -212,8 +204,7 @@ struct Sm70WarpShape { "SM70 Marlin supports CTA_N in {64, 128, 256}."); static_assert(CtaK == 16 || CtaK == 32 || CtaK == 64 || CtaK == 128, "SM70 Marlin supports CTA_K in {16, 32, 64, 128}."); - static_assert(Warps == 4 || Warps == 8, - "SM70 Marlin supports 4 or 8 warps."); + static_assert(Warps == 4 || Warps == 8, "SM70 Marlin supports 4 or 8 warps."); static_assert(WarpM == 32 || WarpM == 64, "SM70 Marlin supports Warp_M in {32, 64}."); static_assert(WarpN == 32 || WarpN == 64, @@ -283,8 +274,7 @@ inline constexpr char const* kSupportedSm70MarlinCtaGeometries = "decomposition, current CUTLASS Volta thread-map support, and " "phase-aware SM70 Marlin warp-K offset support."; -inline bool sm70_marlin_cta_geometry_is_supported( - Sm70CtaGeometry geometry) { +inline bool sm70_marlin_cta_geometry_is_supported(Sm70CtaGeometry geometry) { return sm70_marlin_cta_geometry_is_supported_constexpr( geometry.cta_m, geometry.cta_n, geometry.cta_k, geometry.warps, geometry.warp_m, geometry.warp_n, geometry.warp_k); @@ -300,8 +290,10 @@ inline int sm70_marlin_auto_packed_macro_n(int64_t size_n) { if (size_n % 64 == 0) { return 64; } - TORCH_CHECK(false, "SM70 Marlin requires size_n divisible " - "by 64. Got size_n = ", size_n, "."); + TORCH_CHECK(false, + "SM70 Marlin requires size_n divisible " + "by 64. Got size_n = ", + size_n, "."); return 0; } @@ -309,14 +301,12 @@ inline bool sm70_marlin_packed_macro_n_is_supported(int packed_macro_n) { return packed_macro_n == 64 || packed_macro_n == 128 || packed_macro_n == 256; } -inline bool sm70_marlin_requested_split_k_is_supported( - int requested_split_k) { +inline bool sm70_marlin_requested_split_k_is_supported(int requested_split_k) { return requested_split_k == 1 || requested_split_k == 2 || requested_split_k == 4 || requested_split_k == 8; } -inline Sm70MarlinAutoParams sm70_marlin_default_auto_params( - int packed_macro_n); +inline Sm70MarlinAutoParams sm70_marlin_default_auto_params(int packed_macro_n); inline Sm70CtaGeometry sm70_marlin_default_geometry_from_packed_macro_n( int packed_macro_n) { @@ -328,8 +318,8 @@ inline Sm70CtaGeometry sm70_marlin_default_geometry_from_packed_macro_n( case 256: return {32, 256, 32, 4, 32, 64, 32}; default: - TORCH_CHECK(false, "Unsupported SM70 Marlin packed macro-N=", - packed_macro_n, + TORCH_CHECK(false, + "Unsupported SM70 Marlin packed macro-N=", packed_macro_n, "."); } return {0, 0, 0, 0, 0, 0, 0}; @@ -337,10 +327,8 @@ inline Sm70CtaGeometry sm70_marlin_default_geometry_from_packed_macro_n( inline Sm70MarlinAutoParams sm70_marlin_default_auto_params( int packed_macro_n) { - return {sm70_marlin_default_geometry_from_packed_macro_n(packed_macro_n), - 1, - true, - packed_macro_n}; + return {sm70_marlin_default_geometry_from_packed_macro_n(packed_macro_n), 1, + true, packed_macro_n}; } inline bool sm70_marlin_parse_int_component(char const*& cursor, int& value) { @@ -386,9 +374,10 @@ inline bool sm70_marlin_try_get_geometry_env(char const* env_name, if (value == nullptr || value[0] == '\0') { return false; } - TORCH_CHECK(sm70_marlin_parse_geometry_env_value(value, geometry), - "Invalid ", env_name, " value '", value, - "'. Expected {CTA_M}x{CTA_N}x{CTA_K}x{Warps}x{WarpM}x{WarpN}x{WarpK}."); + TORCH_CHECK( + sm70_marlin_parse_geometry_env_value(value, geometry), "Invalid ", + env_name, " value '", value, + "'. Expected {CTA_M}x{CTA_N}x{CTA_K}x{Warps}x{WarpM}x{WarpN}x{WarpK}."); return true; } @@ -409,8 +398,8 @@ inline bool sm70_marlin_try_get_split_k_env(char const* env_name, return true; } -inline bool sm70_marlin_try_get_metadata_env( - char const* env_name, bool& use_metadata_vector_words) { +inline bool sm70_marlin_try_get_metadata_env(char const* env_name, + bool& use_metadata_vector_words) { char const* value = std::getenv(env_name); if (value == nullptr || value[0] == '\0') { return false; @@ -433,9 +422,9 @@ inline bool sm70_marlin_env_is_set(char const* env_name) { return value != nullptr && value[0] != '\0'; } -inline bool sm70_marlin_any_auto_env_is_set( - char const* geometry_env_name, char const* split_k_env_name, - char const* metadata_env_name) { +inline bool sm70_marlin_any_auto_env_is_set(char const* geometry_env_name, + char const* split_k_env_name, + char const* metadata_env_name) { return sm70_marlin_env_is_set(geometry_env_name) || sm70_marlin_env_is_set(split_k_env_name) || sm70_marlin_env_is_set(metadata_env_name); @@ -447,14 +436,13 @@ inline void validate_sm70_marlin_packed_macro_n(char const* op_name, TORCH_CHECK(sm70_marlin_packed_macro_n_is_supported(packed_macro_n), "Unsupported SM70 Marlin packed macro-N for ", op_name, ": ", packed_macro_n, "."); - TORCH_CHECK(packed_macro_n % geometry.cta_n == 0, - "SM70 Marlin ", op_name, + TORCH_CHECK(packed_macro_n % geometry.cta_n == 0, "SM70 Marlin ", op_name, " requires PackedMacroN divisible by CTA_N. Got PackedMacroN=", packed_macro_n, ", CTA_N=", geometry.cta_n, "."); } -inline void validate_sm70_marlin_auto_params( - char const* op_name, Sm70MarlinAutoParams params) { +inline void validate_sm70_marlin_auto_params(char const* op_name, + Sm70MarlinAutoParams params) { TORCH_CHECK(sm70_marlin_cta_geometry_is_supported(params.geometry), "Unsupported SM70 Marlin CTA geometry for ", op_name, ": ", params.geometry.cta_m, "x", params.geometry.cta_n, "x", @@ -474,8 +462,7 @@ inline Sm70MarlinAutoParams sm70_marlin_auto_params_from_env( char const* op_name, char const* geometry_env_name, char const* split_k_env_name, char const* metadata_env_name, int packed_macro_n) { - Sm70MarlinAutoParams params = - sm70_marlin_default_auto_params(packed_macro_n); + Sm70MarlinAutoParams params = sm70_marlin_default_auto_params(packed_macro_n); sm70_marlin_try_get_geometry_env(geometry_env_name, params.geometry); sm70_marlin_try_get_split_k_env(split_k_env_name, params.requested_split_k); sm70_marlin_try_get_metadata_env(metadata_env_name, @@ -485,31 +472,27 @@ inline Sm70MarlinAutoParams sm70_marlin_auto_params_from_env( } inline bool sm70_marlin_dense_auto_env_is_set() { - return sm70_marlin_any_auto_env_is_set( - "SM70_MARLIN_DENSE_CTA_GEOMETRY", "SM70_MARLIN_DENSE_SPLIT_K", - "SM70_MARLIN_DENSE_METADATA_CACHE"); + return sm70_marlin_any_auto_env_is_set("SM70_MARLIN_DENSE_CTA_GEOMETRY", + "SM70_MARLIN_DENSE_SPLIT_K", + "SM70_MARLIN_DENSE_METADATA_CACHE"); } inline Sm70MarlinAutoParams sm70_marlin_dense_auto_params_from_env( Sm70MarlinDenseAutoParamsContext const& ctx) { return sm70_marlin_auto_params_from_env( - "Dense", "SM70_MARLIN_DENSE_CTA_GEOMETRY", - "SM70_MARLIN_DENSE_SPLIT_K", "SM70_MARLIN_DENSE_METADATA_CACHE", - ctx.packed_macro_n); + "Dense", "SM70_MARLIN_DENSE_CTA_GEOMETRY", "SM70_MARLIN_DENSE_SPLIT_K", + "SM70_MARLIN_DENSE_METADATA_CACHE", ctx.packed_macro_n); } -inline bool -sm70_marlin_dense_try_select_quanttrio_qwen3_6_27b_awq_params( - Sm70MarlinDenseAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool sm70_marlin_dense_try_select_quanttrio_qwen3_6_27b_awq_params( + Sm70MarlinDenseAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4") != 0 || - ctx.group_size != 128 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4") != 0 || ctx.group_size != 128 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -596,16 +579,14 @@ sm70_marlin_dense_try_select_quanttrio_qwen3_6_27b_awq_params( inline bool sm70_marlin_dense_try_select_cyankiwi_qwen3_6_27b_awq_bf16_nvfp4_params( - Sm70MarlinDenseAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinDenseAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "nvfp4") != 0 || - ctx.group_size != 16 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "nvfp4") != 0 || ctx.group_size != 16 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -614,147 +595,144 @@ sm70_marlin_dense_try_select_cyankiwi_qwen3_6_27b_awq_bf16_nvfp4_params( if (ctx.size_n == 2048 && ctx.size_k == 5120) { if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,64,16}, 4, true); + return set_params({32, 64, 64, 4, 32, 64, 16}, 4, true); } if (ctx.size_m <= 64) { - return set_params({64,64,64,4,32,64,32}, 4, true); + return set_params({64, 64, 64, 4, 32, 64, 32}, 4, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, false); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, false); } if (ctx.size_n == 3584 && ctx.size_k == 5120) { if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,64,16}, 4, true); + return set_params({32, 64, 64, 4, 32, 64, 16}, 4, true); } if (ctx.size_m <= 64) { - return set_params({64,64,32,4,32,32,32}, 4, true); + return set_params({64, 64, 32, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 4352 && ctx.size_k == 5120) { if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 64) { - return set_params({64,64,32,4,32,32,32}, 4, false); + return set_params({64, 64, 32, 4, 32, 32, 32}, 4, false); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, false); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, false); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 5120 && ctx.size_k == 768) { if (ctx.size_m <= 1) { - return set_params({64,64,64,4,32,64,32}, 1, false); + return set_params({64, 64, 64, 4, 32, 64, 32}, 1, false); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,64,16}, 1, false); + return set_params({32, 64, 64, 4, 32, 64, 16}, 1, false); } if (ctx.size_m <= 64) { - return set_params({32,64,32,4,32,32,16}, 1, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, false); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, false); } if (ctx.size_n == 5120 && ctx.size_k == 1536) { if (ctx.size_m <= 1) { - return set_params({32,64,64,4,32,64,16}, 1, true); + return set_params({32, 64, 64, 4, 32, 64, 16}, 1, true); } if (ctx.size_m <= 32) { - return set_params({64,64,32,4,32,32,32}, 1, true); + return set_params({64, 64, 32, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 64) { - return set_params({64,64,32,4,32,32,32}, 1, false); + return set_params({64, 64, 32, 4, 32, 32, 32}, 1, false); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 5120 && ctx.size_k == 2176) { if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); } if (ctx.size_m <= 32) { - return set_params({32,64,128,4,32,64,32}, 1, true); + return set_params({32, 64, 128, 4, 32, 64, 32}, 1, true); } if (ctx.size_m <= 64) { - return set_params({32,128,128,8,32,64,32}, 1, true); + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 5120 && ctx.size_k == 4352) { if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); } if (ctx.size_m <= 32) { - return set_params({32,128,64,8,32,32,32}, 2, true); + return set_params({32, 128, 64, 8, 32, 32, 32}, 2, true); } if (ctx.size_m <= 64) { - return set_params({32,128,128,8,32,64,32}, 1, true); + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 8704 && ctx.size_k == 5120) { if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 4, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 32) { - return set_params({32,128,32,4,32,32,32}, 4, false); + return set_params({32, 128, 32, 4, 32, 32, 32}, 4, false); } if (ctx.size_m <= 64) { - return set_params({64,128,128,8,64,64,32}, 1, true); + return set_params({64, 128, 128, 8, 64, 64, 32}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } return false; } -inline bool -sm70_marlin_dense_try_select_nvidia_glm_4_7_nvfp4_params( - Sm70MarlinDenseAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool sm70_marlin_dense_try_select_nvidia_glm_4_7_nvfp4_params( + Sm70MarlinDenseAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "nvfp4") != 0 || - ctx.group_size != 16 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "nvfp4") != 0 || ctx.group_size != 16 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -763,293 +741,282 @@ sm70_marlin_dense_try_select_nvidia_glm_4_7_nvfp4_params( if (ctx.size_n == 384 && ctx.size_k == 5120) { if (ctx.size_m <= 1) { - return set_params({32,64,64,4,32,64,16}, 8, true); + return set_params({32, 64, 64, 4, 32, 64, 16}, 8, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 8, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 8, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,64,16}, 8, true); + return set_params({32, 64, 64, 4, 32, 64, 16}, 8, true); } if (ctx.size_m <= 2048) { - return set_params({64,64,64,4,64,32,32}, 1, true); + return set_params({64, 64, 64, 4, 64, 32, 32}, 1, true); } - return set_params({64,128,32,4,32,64,32}, 1, true); + return set_params({64, 128, 32, 4, 32, 64, 32}, 1, true); } if (ctx.size_n == 768 && ctx.size_k == 5120) { if (ctx.size_m <= 1) { - return set_params({32,64,64,4,32,32,32}, 8, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 8, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,64,16}, 8, false); + return set_params({32, 64, 64, 4, 32, 64, 16}, 8, false); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 2048) { - return set_params({128,64,32,4,32,64,32}, 1, true); + return set_params({128, 64, 32, 4, 32, 64, 32}, 1, true); } - return set_params({64,128,32,4,32,64,32}, 1, true); + return set_params({64, 128, 32, 4, 32, 64, 32}, 1, true); } if (ctx.size_n == 3072 && ctx.size_k == 5120) { if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,64,16}, 4, true); + return set_params({32, 64, 64, 4, 32, 64, 16}, 4, true); } if (ctx.size_m <= 64) { - return set_params({64,64,32,4,32,32,32}, 4, true); + return set_params({64, 64, 32, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 2048) { - return set_params({128,128,64,8,64,64,32}, 1, true); + return set_params({128, 128, 64, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 5120 && ctx.size_k == 192) { if (ctx.size_m <= 1) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 32) { - return set_params({32,64,32,4,32,32,16}, 1, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, true); } if (ctx.size_m <= 64) { - return set_params({32,64,32,4,32,32,16}, 1, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,128,32,4,64,64,32}, 1, true); + return set_params({128, 128, 32, 4, 64, 64, 32}, 1, true); } - return set_params({128,128,32,4,64,64,32}, 1, true); + return set_params({128, 128, 32, 4, 64, 64, 32}, 1, true); } if (ctx.size_n == 5120 && ctx.size_k == 384) { if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 1, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, true); } if (ctx.size_m <= 32) { - return set_params({32,64,32,4,32,32,16}, 1, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, true); } if (ctx.size_m <= 64) { - return set_params({32,64,32,4,32,32,16}, 1, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 5120 && ctx.size_k == 1536) { if (ctx.size_m <= 1) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 64) { - return set_params({32,128,32,4,32,32,32}, 1, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 5120 && ctx.size_k == 3072) { if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 8, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 8, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,64,16}, 2, false); + return set_params({32, 64, 64, 4, 32, 64, 16}, 2, false); } if (ctx.size_m <= 64) { - return set_params({32,128,128,8,32,64,32}, 1, true); + return set_params({32, 128, 128, 8, 32, 64, 32}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 6144 && ctx.size_k == 5120) { if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 4, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 32) { - return set_params({32,128,32,4,32,32,32}, 4, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 64) { - return set_params({64,128,32,4,32,64,32}, 4, false); + return set_params({64, 128, 32, 4, 32, 64, 32}, 4, false); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } return false; } -inline bool -sm70_marlin_dense_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params( - Sm70MarlinDenseAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool sm70_marlin_dense_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params( + Sm70MarlinDenseAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; return true; }; - if (std::strcmp(ctx.quant_format, "fp8_e4m3") == 0 && - ctx.group_size == 512 && + if (std::strcmp(ctx.quant_format, "fp8_e4m3") == 0 && ctx.group_size == 512 && ctx.size_n == 2048 && ctx.size_k == 512) { if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 1, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, true); } if (ctx.size_m <= 32) { - return set_params({32,64,32,4,32,32,16}, 1, false); + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, false); } if (ctx.size_m <= 64) { - return set_params({32,64,32,4,32,32,16}, 1, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (std::strcmp(ctx.quant_format, "fp8_e4m3") == 0 && - ctx.group_size == 1024 && - ctx.size_n == 2048 && ctx.size_k == 1024) { + ctx.group_size == 1024 && ctx.size_n == 2048 && ctx.size_k == 1024) { if (ctx.size_m <= 1) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (std::strcmp(ctx.quant_format, "fp8_e4m3") == 0 && - ctx.group_size == 2048 && - ctx.size_n == 1536 && ctx.size_k == 2048) { + ctx.group_size == 2048 && ctx.size_n == 1536 && ctx.size_k == 2048) { if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 4, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 4, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 2048) { - return set_params({64,128,32,4,32,64,32}, 1, false); + return set_params({64, 128, 32, 4, 32, 64, 32}, 1, false); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (std::strcmp(ctx.quant_format, "fp8_e4m3") == 0 && - ctx.group_size == 2048 && - ctx.size_n == 2560 && ctx.size_k == 2048) { + ctx.group_size == 2048 && ctx.size_n == 2560 && ctx.size_k == 2048) { if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 8, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 8, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 2, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 2, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,64,16}, 1, true); + return set_params({32, 64, 64, 4, 32, 64, 16}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - if (std::strcmp(ctx.quant_format, "nvfp4") == 0 && - ctx.group_size == 16 && + if (std::strcmp(ctx.quant_format, "nvfp4") == 0 && ctx.group_size == 16 && ctx.size_n == 128 && ctx.size_k == 2048) { if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 4, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 4, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 2048) { - return set_params({64,64,64,4,32,64,32}, 1, true); + return set_params({64, 64, 64, 4, 32, 64, 32}, 1, true); } - return set_params({128,64,128,8,64,64,32}, 1, true); + return set_params({128, 64, 128, 8, 64, 64, 32}, 1, true); } - if (std::strcmp(ctx.quant_format, "nvfp4") == 0 && - ctx.group_size == 16 && + if (std::strcmp(ctx.quant_format, "nvfp4") == 0 && ctx.group_size == 16 && ctx.size_n == 256 && ctx.size_k == 2048) { if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 4, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 4, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 2048) { - return set_params({128,64,128,8,64,64,32}, 1, false); + return set_params({128, 64, 128, 8, 64, 64, 32}, 1, false); } - return set_params({128,64,32,4,32,64,32}, 1, false); + return set_params({128, 64, 32, 4, 32, 64, 32}, 1, false); } - if (std::strcmp(ctx.quant_format, "nvfp4") == 0 && - ctx.group_size == 16 && + if (std::strcmp(ctx.quant_format, "nvfp4") == 0 && ctx.group_size == 16 && ctx.size_n == 2048 && ctx.size_k == 64) { if (ctx.size_m <= 1) { - return set_params({32,64,64,4,32,64,16}, 1, false); + return set_params({32, 64, 64, 4, 32, 64, 16}, 1, false); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,64,16}, 1, false); + return set_params({32, 64, 64, 4, 32, 64, 16}, 1, false); } if (ctx.size_m <= 2048) { - return set_params({32,128,32,4,32,32,32}, 1, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } - return set_params({128,64,32,4,32,64,32}, 1, false); + return set_params({128, 64, 32, 4, 32, 64, 32}, 1, false); } - if (std::strcmp(ctx.quant_format, "nvfp4") == 0 && - ctx.group_size == 16 && + if (std::strcmp(ctx.quant_format, "nvfp4") == 0 && ctx.group_size == 16 && ctx.size_n == 2048 && ctx.size_k == 128) { if (ctx.size_m <= 1) { - return set_params({32,64,64,4,32,32,32}, 1, false); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, false); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,64,16}, 1, false); + return set_params({32, 64, 64, 4, 32, 64, 16}, 1, false); } if (ctx.size_m <= 2048) { - return set_params({128,64,32,4,32,64,32}, 1, true); + return set_params({128, 64, 32, 4, 32, 64, 32}, 1, true); } - return set_params({128,64,32,4,32,64,32}, 1, true); + return set_params({128, 64, 32, 4, 32, 64, 32}, 1, true); } return false; @@ -1057,16 +1024,14 @@ sm70_marlin_dense_try_select_nvidia_qwen3_6_35b_a3b_nvfp4_params( inline bool sm70_marlin_dense_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_params( - Sm70MarlinDenseAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { + Sm70MarlinDenseAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "nvfp4") != 0 || - ctx.group_size != 16 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "nvfp4") != 0 || ctx.group_size != 16 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -1075,115 +1040,112 @@ sm70_marlin_dense_try_select_nvidia_qwen3_next_80b_a3b_thinking_nvfp4_params( if (ctx.size_n == 128 && ctx.size_k == 2048) { if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 4, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 4, true); } if (ctx.size_m <= 32) { - return set_params({32,64,32,4,32,32,16}, 4, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 4, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 2048) { - return set_params({64,64,64,4,32,64,32}, 1, true); + return set_params({64, 64, 64, 4, 32, 64, 32}, 1, true); } - return set_params({128,64,128,8,64,64,32}, 1, true); + return set_params({128, 64, 128, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 256 && ctx.size_k == 2048) { if (ctx.size_m <= 1) { - return set_params({32,64,32,4,32,32,16}, 4, true); + return set_params({32, 64, 32, 4, 32, 32, 16}, 4, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,32,32}, 4, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 4, true); } if (ctx.size_m <= 2048) { - return set_params({128,64,128,8,64,64,32}, 1, true); + return set_params({128, 64, 128, 8, 64, 64, 32}, 1, true); } - return set_params({128,64,32,4,32,64,32}, 1, true); + return set_params({128, 64, 32, 4, 32, 64, 32}, 1, true); } if (ctx.size_n == 2048 && ctx.size_k == 64) { if (ctx.size_m <= 1) { - return set_params({32,64,64,4,32,64,16}, 1, true); + return set_params({32, 64, 64, 4, 32, 64, 16}, 1, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,32,32}, 1, false); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, false); } if (ctx.size_m <= 2048) { - return set_params({64,64,32,4,32,32,32}, 1, true); + return set_params({64, 64, 32, 4, 32, 32, 32}, 1, true); } - return set_params({128,64,32,4,32,64,32}, 1, true); + return set_params({128, 64, 32, 4, 32, 64, 32}, 1, true); } if (ctx.size_n == 2048 && ctx.size_k == 128) { if (ctx.size_m <= 1) { - return set_params({32,128,32,4,32,32,32}, 1, true); + return set_params({32, 128, 32, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,64,16}, 1, true); + return set_params({32, 64, 64, 4, 32, 64, 16}, 1, true); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,64,32,4,32,64,32}, 1, true); + return set_params({128, 64, 32, 4, 32, 64, 32}, 1, true); } - return set_params({128,64,32,4,32,64,32}, 1, false); + return set_params({128, 64, 32, 4, 32, 64, 32}, 1, false); } if (ctx.size_n == 2048 && ctx.size_k == 512) { if (ctx.size_m <= 1) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 1, false); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, false); } if (ctx.size_m <= 64) { - return set_params({32,64,32,4,32,32,16}, 1, false); + return set_params({32, 64, 32, 4, 32, 32, 16}, 1, false); } if (ctx.size_m <= 2048) { - return set_params({128,128,32,4,64,64,32}, 1, true); + return set_params({128, 128, 32, 4, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } if (ctx.size_n == 2048 && ctx.size_k == 1024) { if (ctx.size_m <= 1) { - return set_params({64,64,32,4,32,32,32}, 1, true); + return set_params({64, 64, 32, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 32) { - return set_params({32,64,64,4,32,32,32}, 1, false); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, false); } if (ctx.size_m <= 64) { - return set_params({32,64,64,4,32,32,32}, 1, true); + return set_params({32, 64, 64, 4, 32, 32, 32}, 1, true); } if (ctx.size_m <= 2048) { - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } - return set_params({128,256,32,8,64,64,32}, 1, true); + return set_params({128, 256, 32, 8, 64, 64, 32}, 1, true); } return false; } -inline bool -sm70_marlin_dense_try_select_cyankiwi_minimax_m2_7_awq_4bit_params( - Sm70MarlinDenseAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool sm70_marlin_dense_try_select_cyankiwi_minimax_m2_7_awq_4bit_params( + Sm70MarlinDenseAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4b8") != 0 || - ctx.group_size != 32 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4b8") != 0 || ctx.group_size != 32 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -1246,18 +1208,15 @@ sm70_marlin_dense_try_select_cyankiwi_minimax_m2_7_awq_4bit_params( return false; } -inline bool -sm70_marlin_dense_try_select_cyankiwi_glm_4_7_awq_4bit_params( - Sm70MarlinDenseAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool sm70_marlin_dense_try_select_cyankiwi_glm_4_7_awq_4bit_params( + Sm70MarlinDenseAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4b8") != 0 || - ctx.group_size != 32 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4b8") != 0 || ctx.group_size != 32 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -1323,18 +1282,15 @@ sm70_marlin_dense_try_select_cyankiwi_glm_4_7_awq_4bit_params( return false; } -inline bool -sm70_marlin_dense_try_select_quanttrio_glm_4_7_awq_params( - Sm70MarlinDenseAutoParamsContext const& ctx, - Sm70MarlinAutoParams& params) { +inline bool sm70_marlin_dense_try_select_quanttrio_glm_4_7_awq_params( + Sm70MarlinDenseAutoParamsContext const& ctx, Sm70MarlinAutoParams& params) { if (ctx.quant_format == nullptr || - std::strcmp(ctx.quant_format, "uint4") != 0 || - ctx.group_size != 128 || ctx.size_m <= 0) { + std::strcmp(ctx.quant_format, "uint4") != 0 || ctx.group_size != 128 || + ctx.size_m <= 0) { return false; } - auto const set_params = [&](Sm70CtaGeometry geometry, - int requested_split_k, + auto const set_params = [&](Sm70CtaGeometry geometry, int requested_split_k, bool use_metadata_vector_words) { params = {geometry, requested_split_k, use_metadata_vector_words, ctx.packed_macro_n}; @@ -1426,16 +1382,16 @@ inline Sm70MarlinAutoParams sm70_marlin_dense_auto_params( char const* quant_format, int64_t group_size, int64_t size_m, int64_t size_n, int64_t size_k) { Sm70MarlinDenseAutoParamsContext const ctx{ - quant_format, group_size, size_m, size_n, size_k, - sm70_marlin_auto_packed_macro_n(size_n)}; + quant_format, group_size, size_m, + size_n, size_k, sm70_marlin_auto_packed_macro_n(size_n)}; if (sm70_marlin_dense_auto_env_is_set()) { return sm70_marlin_dense_auto_params_from_env(ctx); } Sm70MarlinAutoParams params{}; - if (sm70_marlin_dense_try_select_quanttrio_qwen3_6_27b_awq_params( - ctx, params)) { + if (sm70_marlin_dense_try_select_quanttrio_qwen3_6_27b_awq_params(ctx, + params)) { validate_sm70_marlin_auto_params("Dense", params); return params; } @@ -1444,8 +1400,7 @@ inline Sm70MarlinAutoParams sm70_marlin_dense_auto_params( validate_sm70_marlin_auto_params("Dense", params); return params; } - if (sm70_marlin_dense_try_select_nvidia_glm_4_7_nvfp4_params( - ctx, params)) { + if (sm70_marlin_dense_try_select_nvidia_glm_4_7_nvfp4_params(ctx, params)) { validate_sm70_marlin_auto_params("Dense", params); return params; } @@ -1464,13 +1419,12 @@ inline Sm70MarlinAutoParams sm70_marlin_dense_auto_params( validate_sm70_marlin_auto_params("Dense", params); return params; } - if (sm70_marlin_dense_try_select_cyankiwi_glm_4_7_awq_4bit_params( - ctx, params)) { + if (sm70_marlin_dense_try_select_cyankiwi_glm_4_7_awq_4bit_params(ctx, + params)) { validate_sm70_marlin_auto_params("Dense", params); return params; } - if (sm70_marlin_dense_try_select_quanttrio_glm_4_7_awq_params( - ctx, params)) { + if (sm70_marlin_dense_try_select_quanttrio_glm_4_7_awq_params(ctx, params)) { validate_sm70_marlin_auto_params("Dense", params); return params; } @@ -1478,27 +1432,25 @@ inline Sm70MarlinAutoParams sm70_marlin_dense_auto_params( return sm70_marlin_default_auto_params(ctx.packed_macro_n); } -inline void validate_sm70_marlin_dense_cta_geometry_supported(char const* op_name, - Sm70CtaGeometry geometry) { +inline void validate_sm70_marlin_dense_cta_geometry_supported( + char const* op_name, Sm70CtaGeometry geometry) { TORCH_CHECK(sm70_marlin_cta_geometry_is_supported(geometry), "Unsupported SM70 Marlin CTA geometry for ", op_name, ": ", - geometry.cta_m, "x", geometry.cta_n, "x", geometry.cta_k, - "x", geometry.warps, "x", geometry.warp_m, "x", - geometry.warp_n, "x", geometry.warp_k, - ". Supported geometries are ", + geometry.cta_m, "x", geometry.cta_n, "x", geometry.cta_k, "x", + geometry.warps, "x", geometry.warp_m, "x", geometry.warp_n, "x", + geometry.warp_k, ". Supported geometries are ", kSupportedSm70MarlinCtaGeometries, "."); } inline void validate_sm70_marlin_dense_cta_n_alignment(char const* op_name, Sm70CtaGeometry geometry, int64_t size_n) { - TORCH_CHECK( - size_n % geometry.cta_n == 0 && size_n % kQuantTileN == 0, - "SM70 Marlin requires size_n divisible by both CTA_N and 64 for ", - op_name, " with CTA geometry ", geometry.cta_m, "x", geometry.cta_n, - "x", geometry.cta_k, "x", geometry.warps, "x", geometry.warp_m, - "x", geometry.warp_n, "x", geometry.warp_k, ". Got size_n = ", - size_n, "."); + TORCH_CHECK(size_n % geometry.cta_n == 0 && size_n % kQuantTileN == 0, + "SM70 Marlin requires size_n divisible by both CTA_N and 64 for ", + op_name, " with CTA geometry ", geometry.cta_m, "x", + geometry.cta_n, "x", geometry.cta_k, "x", geometry.warps, "x", + geometry.warp_m, "x", geometry.warp_n, "x", geometry.warp_k, + ". Got size_n = ", size_n, "."); } } // namespace marlin::sm70 diff --git a/csrc/quantization/marlin/sm70_marlin_dispatch.cu b/csrc/quantization/marlin/sm70_marlin_dispatch.cu index fe7fa0f408..51aa145884 100644 --- a/csrc/quantization/marlin/sm70_marlin_dispatch.cu +++ b/csrc/quantization/marlin/sm70_marlin_dispatch.cu @@ -43,15 +43,13 @@ __device__ __forceinline__ int sm70_inverse_zp_interleave(int idx, return inv[idx]; } -__device__ __forceinline__ int sm70_inverse_scale_perm(int idx, - int perm_len) { +__device__ __forceinline__ int sm70_inverse_scale_perm(int idx, int perm_len) { if (perm_len == 64) { return (idx % 8) * 8 + idx / 8; } - constexpr int inv[32] = {0, 1, 8, 9, 16, 17, 24, 25, - 2, 3, 10, 11, 18, 19, 26, 27, - 4, 5, 12, 13, 20, 21, 28, 29, - 6, 7, 14, 15, 22, 23, 30, 31}; + constexpr int inv[32] = {0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, + 11, 18, 19, 26, 27, 4, 5, 12, 13, 20, 21, + 28, 29, 6, 7, 14, 15, 22, 23, 30, 31}; return inv[idx]; } @@ -66,10 +64,9 @@ __global__ void sm70_unpermute_half_metadata_kernel( } __global__ void sm70_unpack_dense_zp_to_half_kernel( - const int32_t* __restrict__ packed_zp, - const half* __restrict__ scales, half* __restrict__ out, - int64_t total, int64_t size_n, int64_t packed_n, int num_bits, - bool is_a_8bit) { + const int32_t* __restrict__ packed_zp, const half* __restrict__ scales, + half* __restrict__ out, int64_t total, int64_t size_n, int64_t packed_n, + int num_bits, bool is_a_8bit) { const int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= total) return; @@ -81,21 +78,18 @@ __global__ void sm70_unpack_dense_zp_to_half_kernel( if (!is_a_8bit) { packed_col_in_block = (packed_col_in_block / pack_factor) * pack_factor + - sm70_inverse_zp_interleave(packed_col_in_block % pack_factor, - num_bits); + sm70_inverse_zp_interleave(packed_col_in_block % pack_factor, num_bits); } const int64_t packed_col = block64 * 64 + packed_col_in_block; - const uint32_t word = reinterpret_cast(packed_zp) - [(idx / size_n) * packed_n + packed_col / pack_factor]; - const uint32_t zp = - (word >> ((packed_col % pack_factor) * num_bits)) & - ((1u << num_bits) - 1u); + const uint32_t word = reinterpret_cast( + packed_zp)[(idx / size_n) * packed_n + packed_col / pack_factor]; + const uint32_t zp = (word >> ((packed_col % pack_factor) * num_bits)) & + ((1u << num_bits) - 1u); out[idx] = __float2half(static_cast(zp) * __half2float(scales[idx])); } torch::Tensor sm70_dense_scales_as_logical(torch::Tensor const& b_scales, - int64_t size_k, - int64_t group_size, + int64_t size_k, int64_t group_size, bool is_a_8bit) { TORCH_CHECK(b_scales.scalar_type() == at::ScalarType::Half, "SM70 Marlin adapter expects fp16 scales."); @@ -103,15 +97,14 @@ torch::Tensor sm70_dense_scales_as_logical(torch::Tensor const& b_scales, group_size < size_k && group_size != -1 && !is_a_8bit; const int perm_len = uses_group_perm ? 64 : 32; TORCH_CHECK(b_scales.numel() % perm_len == 0, - "SM70 Marlin scale metadata size is not divisible by ", - perm_len); + "SM70 Marlin scale metadata size is not divisible by ", perm_len); auto out = torch::empty_like(b_scales); const int64_t total = out.numel(); constexpr int threads = 256; const int blocks = static_cast((total + threads - 1) / threads); - sm70_unpermute_half_metadata_kernel<<< - blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>( + sm70_unpermute_half_metadata_kernel<<>>( reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(out.data_ptr()), total, perm_len); C10_CUDA_KERNEL_LAUNCH_CHECK(); @@ -143,8 +136,8 @@ torch::Tensor sm70_dense_zp_as_half(torch::Tensor const& b_zeros, const int64_t total = out.numel(); constexpr int threads = 256; const int blocks = static_cast((total + threads - 1) / threads); - sm70_unpack_dense_zp_to_half_kernel<<< - blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>( + sm70_unpack_dense_zp_to_half_kernel<<>>( b_zeros.data_ptr(), reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(out.data_ptr()), total, size_n, @@ -202,8 +195,7 @@ torch::Tensor marlin_gemm( std::optional const& global_scale_or_none, std::optional const& b_zeros_or_none, std::optional const& g_idx_or_none, - std::optional const& perm_or_none, - torch::Tensor& workspace, + std::optional const& perm_or_none, torch::Tensor& workspace, vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { @@ -273,8 +265,7 @@ torch::Tensor marlin_gemm( "SM70 build only supports float16 outputs."); TORCH_CHECK(s_type == vllm::kFloat16 || (b_type == vllm::kFE2M1f && - (s_type == vllm::kFE4M3fn || - s_type == vllm::kFE8M0fnu)), + (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu)), "SM70 build only supports float16 scales, except FP4 uses " "float8_e4m3fn scales for NVFP4 or float8_e8m0fnu scales for " "MXFP4."); @@ -325,10 +316,8 @@ torch::Tensor marlin_gemm( TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); TORCH_CHECK(b_scales.scalar_type() == at::ScalarType::Half || (b_type == vllm::kFE2M1f && - (b_scales.scalar_type() == - at::ScalarType::Float8_e4m3fn || - b_scales.scalar_type() == - at::ScalarType::Float8_e8m0fnu)), + (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn || + b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu)), "SM70 build only supports float16 scales, except FP4 uses " "float8_e4m3fn scales for NVFP4 or float8_e8m0fnu scales for " "MXFP4."); @@ -421,21 +410,18 @@ torch::Tensor marlin_gemm( torch::Tensor b_scales_kernel = b_scales; if (b_scales.scalar_type() == at::ScalarType::Half) { - b_scales_kernel = - sm70_dense_scales_as_logical(b_scales, size_k, group_size, - a_type.size_bits() == 8); + b_scales_kernel = sm70_dense_scales_as_logical(b_scales, size_k, group_size, + a_type.size_bits() == 8); } torch::Tensor global_scale; if (global_scale_or_none.has_value()) { global_scale = global_scale_or_none.value(); - TORCH_CHECK( - b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, - "SM70 Marlin supports global_scale only for " - "nvfp4 format."); + TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, + "SM70 Marlin supports global_scale only for " + "nvfp4 format."); TORCH_CHECK(global_scale.device().is_cuda(), "global_scale is not on GPU"); - TORCH_CHECK(global_scale.is_contiguous(), - "global_scale is not contiguous"); + TORCH_CHECK(global_scale.is_contiguous(), "global_scale is not contiguous"); TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float, "SM70 Marlin NVFP4 expects fp32 global_scale."); TORCH_CHECK(global_scale.numel() == 1, @@ -463,8 +449,7 @@ torch::Tensor marlin_gemm( if (b_zeros_or_none.has_value()) { b_zeros = b_zeros_or_none.value(); TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU"); - TORCH_CHECK(b_zeros.is_contiguous(), - "b_zeros is not contiguous"); + TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous"); } else { b_zeros = torch::empty({0}, options); } @@ -477,9 +462,9 @@ torch::Tensor marlin_gemm( if (has_zp) { if (!is_zp_float && b_zeros.scalar_type() == at::ScalarType::Int) { - b_zeros = sm70_dense_zp_as_half( - b_zeros, b_scales_kernel, size_n, b_type == vllm::kU4 ? 4 : 8, - a_type.size_bits() == 8); + b_zeros = sm70_dense_zp_as_half(b_zeros, b_scales_kernel, size_n, + b_type == vllm::kU4 ? 4 : 8, + a_type.size_bits() == 8); zp_is_float = true; } } @@ -493,8 +478,7 @@ torch::Tensor marlin_gemm( TORCH_CHECK(b_zeros.scalar_type() == at::ScalarType::Half, "SM70 Marlin uint4/uint8 dense path expects fp16 " "zero points."); - TORCH_CHECK(b_zeros.size(1) == size_n, - "b_zeros dim 1 = ", b_zeros.size(1), + TORCH_CHECK(b_zeros.size(1) == size_n, "b_zeros dim 1 = ", b_zeros.size(1), " is not size_n = ", size_n); TORCH_CHECK(num_groups == b_zeros.size(0), "b_zeros dim 0 = ", b_zeros.size(0), @@ -523,30 +507,26 @@ torch::Tensor marlin_gemm( "SM70 Marlin expects int32 packed weights."); TORCH_CHECK(!has_act_order, "act_order is not supported for the SM70 Marlin path."); - TORCH_CHECK(!has_bias, - "SM70 Marlin does not support bias. TODO: add epilogue bias fusion."); + TORCH_CHECK( + !has_bias, + "SM70 Marlin does not support bias. TODO: add epilogue bias fusion."); TORCH_CHECK((b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) || global_scale.numel() == 0, "SM70 Marlin supports global_scale only for " "nvfp4 format."); TORCH_CHECK(!use_atomic_add, "SM70 Marlin does not support atomic-add output."); - TORCH_CHECK(is_k_full, - "SM70 Marlin requires full-K non act-order inputs."); - TORCH_CHECK(size_k % 32 == 0, - "SM70 Marlin requires size_k % 32 == 0."); - TORCH_CHECK(size_n % 64 == 0, - "SM70 Marlin requires size_n % 64 == 0."); + TORCH_CHECK(is_k_full, "SM70 Marlin requires full-K non act-order inputs."); + TORCH_CHECK(size_k % 32 == 0, "SM70 Marlin requires size_k % 32 == 0."); + TORCH_CHECK(size_n % 64 == 0, "SM70 Marlin requires size_n % 64 == 0."); TORCH_CHECK(group_size == -1 || group_size > 0, - "SM70 Marlin received invalid group_size = ", - group_size); + "SM70 Marlin received invalid group_size = ", group_size); if (b_type == vllm::kU4) { TORCH_CHECK(size_k % 32 == 0, "SM70 Marlin uint4 dense path requires size_k % 32 == 0."); - TORCH_CHECK( - has_zp && zp_is_float, - "SM70 Marlin uint4 dense path requires fp16 zero points."); + TORCH_CHECK(has_zp && zp_is_float, + "SM70 Marlin uint4 dense path requires fp16 zero points."); return sm70_marlin_u4_gemm(a, c, b_q_weight, b_scales_kernel, b_zeros, size_m, size_n, size_k, group_size); } @@ -554,17 +534,15 @@ torch::Tensor marlin_gemm( if (b_type == vllm::kU8) { TORCH_CHECK(size_k % 32 == 0, "SM70 Marlin uint8 dense path requires size_k % 32 == 0."); - TORCH_CHECK( - has_zp && zp_is_float, - "SM70 Marlin uint8 dense path requires fp16 zero points."); + TORCH_CHECK(has_zp && zp_is_float, + "SM70 Marlin uint8 dense path requires fp16 zero points."); return sm70_marlin_u8_gemm(a, c, b_q_weight, b_scales_kernel, b_zeros, size_m, size_n, size_k, group_size); } if (b_type == vllm::kU8B128) { - TORCH_CHECK( - size_k % 32 == 0, - "SM70 Marlin uint8b128 dense path requires size_k % 32 == 0."); + TORCH_CHECK(size_k % 32 == 0, + "SM70 Marlin uint8b128 dense path requires size_k % 32 == 0."); TORCH_CHECK(!has_zp && !is_zp_float, "SM70 Marlin uint8b128 does not support " "zero-point metadata."); @@ -573,9 +551,8 @@ torch::Tensor marlin_gemm( } if (b_type == vllm::kFE4M3fn) { - TORCH_CHECK( - size_k % 32 == 0, - "SM70 Marlin FP8 dense path requires size_k % 32 == 0."); + TORCH_CHECK(size_k % 32 == 0, + "SM70 Marlin FP8 dense path requires size_k % 32 == 0."); TORCH_CHECK(group_size == -1 || group_size == 128, "SM70 Marlin FP8 supports only group_size -1 or " "128. Got ", @@ -588,15 +565,15 @@ torch::Tensor marlin_gemm( } if (b_type == vllm::kFE2M1f) { - TORCH_CHECK( - size_k % 32 == 0, - "SM70 Marlin FP4 dense path requires size_k % 32 == 0."); + TORCH_CHECK(size_k % 32 == 0, + "SM70 Marlin FP4 dense path requires size_k % 32 == 0."); TORCH_CHECK(!has_zp && !is_zp_float, "SM70 Marlin FP4 does not support zero-point " "metadata."); if (s_type == vllm::kFE4M3fn) { - TORCH_CHECK(global_scale.numel() == 1, - "the global_scale parameter must be passed for nvfp4 format."); + TORCH_CHECK( + global_scale.numel() == 1, + "the global_scale parameter must be passed for nvfp4 format."); TORCH_CHECK(group_size == 16, "SM70 Marlin NVFP4 supports only group_size 16. " "Got ", diff --git a/csrc/quantization/marlin/sm70_marlin_fp8_gemm.cu b/csrc/quantization/marlin/sm70_marlin_fp8_gemm.cu index 9520f9546a..4dcc8480ad 100644 --- a/csrc/quantization/marlin/sm70_marlin_fp8_gemm.cu +++ b/csrc/quantization/marlin/sm70_marlin_fp8_gemm.cu @@ -20,26 +20,26 @@ #include "cutlass/layout/tensor_op_multiplicand_sm70.h" #include "cutlass/transform/threadblock/predicated_tile_iterator.h" -using marlin::sm70::Sm70CtaGeometry; -using marlin::sm70::Sm70AtomicFp16Epilogue; -using marlin::sm70::Sm70AtomicFp32Epilogue; -using marlin::sm70::sm70_marlin_narrow_f32_to_f16; -using marlin::sm70::Sm70MarlinGemmTraits; -using marlin::sm70::Sm70SplitKPartition; -using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; -using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; using marlin::sm70::configure_sm70_dynamic_smem; using marlin::sm70::dispatch_sm70_marlin_fp8_geometry; using marlin::sm70::kQuantTileK; using marlin::sm70::kQuantTileN; -using marlin::sm70::sm70_marlin_dense_auto_params; using marlin::sm70::load_qword_vector; using marlin::sm70::qword_from_vector; -using marlin::sm70::sm70_marlin_cta_grid; using marlin::sm70::sm70_active_split_k; +using marlin::sm70::sm70_marlin_cta_grid; +using marlin::sm70::sm70_marlin_dense_auto_params; +using marlin::sm70::sm70_marlin_narrow_f32_to_f16; using marlin::sm70::sm70_splitk_partition; +using marlin::sm70::Sm70AtomicFp16Epilogue; +using marlin::sm70::Sm70AtomicFp32Epilogue; +using marlin::sm70::Sm70CtaGeometry; +using marlin::sm70::Sm70MarlinGemmTraits; +using marlin::sm70::Sm70SplitKPartition; using marlin::sm70::u8_packed_macro_n_qweight_offset_from_logical; using marlin::sm70::u8_packed_macro_n_qweight_word_stride; +using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; +using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; namespace { @@ -55,12 +55,11 @@ class Sm70Fp8IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + using Fragment = cutlass::Array; static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, "SM70 Marlin FP8 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, + static_assert(ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, "SM70 Marlin FP8 IteratorB expects one contiguous iteration " "per 64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, @@ -161,13 +160,12 @@ class Sm70Fp8IteratorB { CUTLASS_DEVICE static int qweight_offset_from_logical(Params const& params, int logical_k, int logical_n) { - return u8_packed_macro_n_qweight_offset_from_logical(params.size_n, logical_k, - logical_n); + return u8_packed_macro_n_qweight_offset_from_logical( + params.size_n, logical_k, logical_n); } CUTLASS_DEVICE - static int qweight_word_stride_from_logical(Params const&, - int logical_n) { + static int qweight_word_stride_from_logical(Params const&, int logical_n) { return u8_packed_macro_n_qweight_word_stride(); } @@ -198,9 +196,8 @@ class Sm70Fp8IteratorB { void cache_current_group_metadata(int group) const { CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; if constexpr (kUseMetadataVectorWords) { cache_metadata_vector_words(c, group, cache_n); } else { @@ -261,8 +258,9 @@ class Sm70Fp8IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin FP8 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin FP8 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_offset_); uint32_t const qword1 = load_qword_vector<1>( @@ -284,19 +282,17 @@ class Sm70Fp8IteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; if constexpr (kGroupSize != -1) { cache_current_group_metadata(scale_group(logical_k_s)); } int const qweight_base_s = qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - uint4 const qwords0 = - load_qword_vector<4>(qweight_ + qweight_base_s); - uint4 const qwords1 = load_qword_vector<4>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint4 const qwords0 = load_qword_vector<4>(qweight_ + qweight_base_s); + uint4 const qwords1 = load_qword_vector<4>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -318,10 +314,9 @@ class Sm70Fp8IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - uint2 const qwords0 = - load_qword_vector<2>(qweight_ + qweight_base_s); - uint2 const qwords1 = load_qword_vector<2>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint2 const qwords0 = load_qword_vector<2>(qweight_ + qweight_base_s); + uint2 const qwords1 = load_qword_vector<2>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -343,8 +338,9 @@ class Sm70Fp8IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin FP8 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin FP8 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_s); uint32_t const qword1 = load_qword_vector<1>( @@ -374,8 +370,7 @@ class Sm70Fp8IteratorB { return; } - if constexpr (kGroupSize != -1 && - ThreadMap::Iterations::kStrided == 1) { + if constexpr (kGroupSize != -1 && ThreadMap::Iterations::kStrided == 1) { int const first_logical_k = k_offset_ + thread_offset_.strided(); cache_current_group_metadata(scale_group(first_logical_k)); } @@ -387,32 +382,29 @@ class Sm70Fp8IteratorB { template struct Sm70Fp8GemmSpec { template - using IteratorB = - Sm70Fp8IteratorB; + using IteratorB = Sm70Fp8IteratorB; }; -template using Sm70Fp8GemmTraits = - Sm70MarlinGemmTraits, CtaM, - CtaN, CtaK, Warps, WarpM, WarpN, WarpK, - GroupSize, PackedMacroN>; + Sm70MarlinGemmTraits, CtaM, CtaN, + CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>; -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_fp8_gemm_kernel( +__global__ __launch_bounds__(Warps * 32, 1) void sm70_marlin_fp8_gemm_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, cutlass::half_t const* __restrict__ b_scales, cutlass::half_t* __restrict__ c, int m, int n, int k, int lda) { using Traits = - Sm70Fp8GemmTraits; + Sm70Fp8GemmTraits; using Mma = typename Traits::Mma; using Epilogue = typename Traits::Epilogue; @@ -432,10 +424,10 @@ void sm70_marlin_fp8_gemm_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -460,18 +452,19 @@ void sm70_marlin_fp8_gemm_kernel( epilogue(output_op, iterator_D, accumulators, iterator_C); } -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_fp8_gemm_splitk_kernel( +__global__ +__launch_bounds__(Warps * 32, 1) void sm70_marlin_fp8_gemm_splitk_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, cutlass::half_t const* __restrict__ b_scales, - cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, int k, int lda, int requested_split_k) { - using Traits = Sm70Fp8GemmTraits; + cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, + int k, int lda, int requested_split_k) { + using Traits = + Sm70Fp8GemmTraits; using Mma = typename Traits::Mma; using AtomicEpilogue = Sm70AtomicFp32Epilogue; @@ -482,14 +475,13 @@ void sm70_marlin_fp8_gemm_splitk_kernel( int const thread_idx = threadIdx.x; int const warp_idx = cutlass::canonical_warp_idx_sync(); int const lane_idx = threadIdx.x % 32; - Sm70SplitKPartition const partition = - sm70_splitk_partition(k, requested_split_k, int(blockIdx.z)); + Sm70SplitKPartition const partition = sm70_splitk_partition( + k, requested_split_k, int(blockIdx.z)); if (partition.partition_k == 0) { return; } - cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, - partition.k_begin}; + cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, partition.k_begin}; cutlass::MatrixCoord tb_offset_B{partition.k_begin, int(blockIdx.y) * CtaN}; cutlass::MatrixCoord tb_offset_C{int(blockIdx.x) * CtaM, int(blockIdx.y) * CtaN}; @@ -497,10 +489,10 @@ void sm70_marlin_fp8_gemm_splitk_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -514,9 +506,8 @@ void sm70_marlin_fp8_gemm_splitk_kernel( mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); typename AtomicEpilogue::OutputTileIterator iterator_D( - typename AtomicEpilogue::OutputTileIterator::Params(layout_c), - c, cutlass::MatrixCoord(m, n), - thread_idx, tb_offset_C); + typename AtomicEpilogue::OutputTileIterator::Params(layout_c), c, + cutlass::MatrixCoord(m, n), thread_idx, tb_offset_C); AtomicEpilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx); @@ -525,20 +516,23 @@ void sm70_marlin_fp8_gemm_splitk_kernel( } // namespace -template -torch::Tensor launch_sm70_marlin_fp8_gemm( - torch::Tensor& a, torch::Tensor& c, torch::Tensor& b_q_weight, - torch::Tensor& b_scales, int64_t size_m, int64_t size_n, int64_t size_k, - int requested_split_k) { +torch::Tensor launch_sm70_marlin_fp8_gemm(torch::Tensor& a, torch::Tensor& c, + torch::Tensor& b_q_weight, + torch::Tensor& b_scales, + int64_t size_m, int64_t size_n, + int64_t size_k, + int requested_split_k) { auto kernel = - sm70_marlin_fp8_gemm_kernel; - using SharedStorage = typename Sm70Fp8GemmTraits< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN, UseMetadataVectorWords>::SharedStorage; + using SharedStorage = + typename Sm70Fp8GemmTraits::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -558,14 +552,13 @@ torch::Tensor launch_sm70_marlin_fp8_gemm( } TORCH_CHECK(size_k % int64_t(CtaK) == 0, - "SM70 Marlin fp8_e4m3 requires K divisible by CTA_K=", - CtaK, " for requested_split_k > 1. Got K=", size_k, + "SM70 Marlin fp8_e4m3 requires K divisible by CTA_K=", CtaK, + " for requested_split_k > 1. Got K=", size_k, ", requested_split_k=", requested_split_k, "."); auto split_kernel = - sm70_marlin_fp8_gemm_splitk_kernel; smem_bytes = configure_sm70_dynamic_smem(split_kernel); @@ -573,8 +566,7 @@ torch::Tensor launch_sm70_marlin_fp8_gemm( // 0009: fp32 split-K scratch (zeroed for the fp32 atomicAdd), M*N floats. // Only allocated on the split_k>1 path, which the auto-selector uses ONLY at // low batch (M<128), so the scratch is bounded and small. - auto c32 = torch::zeros({size_m, size_n}, - a.options().dtype(at::kFloat)); + auto c32 = torch::zeros({size_m, size_n}, a.options().dtype(at::kFloat)); dim3 grid = sm70_marlin_cta_grid(size_m, size_n, CtaM, CtaN); int const active_split_k = @@ -585,9 +577,9 @@ torch::Tensor launch_sm70_marlin_fp8_gemm( reinterpret_cast(b_q_weight.data_ptr()), reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(c.data_ptr()), - c32.data_ptr(), - static_cast(size_m), static_cast(size_n), - static_cast(size_k), static_cast(a.stride(0)), requested_split_k); + c32.data_ptr(), static_cast(size_m), static_cast(size_n), + static_cast(size_k), static_cast(a.stride(0)), + requested_split_k); C10_CUDA_KERNEL_LAUNCH_CHECK(); // 0009: single deterministic-width fp32->fp16 narrow of the scratch. @@ -595,8 +587,8 @@ torch::Tensor launch_sm70_marlin_fp8_gemm( int const narrow_grid = static_cast((numel + narrow_block - 1) / narrow_block); sm70_marlin_narrow_f32_to_f16<256><<>>( - c32.data_ptr(), - reinterpret_cast(c.data_ptr()), numel); + c32.data_ptr(), reinterpret_cast(c.data_ptr()), + numel); C10_CUDA_KERNEL_LAUNCH_CHECK(); return c; @@ -613,8 +605,8 @@ struct Sm70Fp8Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { return launch_sm70_marlin_fp8_gemm const launcher{ - a, c, b_q_weight, b_scales, size_m, size_n, size_k, - params.requested_split_k}; - return dispatch_sm70_marlin_fp8_geometry( - launcher, geometry, params.packed_macro_n, group_size); + a, c, b_q_weight, b_scales, + size_m, size_n, size_k, params.requested_split_k}; + return dispatch_sm70_marlin_fp8_geometry(launcher, geometry, + params.packed_macro_n, group_size); } Sm70Fp8Launcher const launcher{ - a, c, b_q_weight, b_scales, size_m, size_n, size_k, - params.requested_split_k}; - return dispatch_sm70_marlin_fp8_geometry( - launcher, geometry, params.packed_macro_n, group_size); + a, c, b_q_weight, b_scales, + size_m, size_n, size_k, params.requested_split_k}; + return dispatch_sm70_marlin_fp8_geometry(launcher, geometry, + params.packed_macro_n, group_size); } diff --git a/csrc/quantization/marlin/sm70_marlin_gemm.cuh b/csrc/quantization/marlin/sm70_marlin_gemm.cuh index 426fac1c0f..a968a54db6 100644 --- a/csrc/quantization/marlin/sm70_marlin_gemm.cuh +++ b/csrc/quantization/marlin/sm70_marlin_gemm.cuh @@ -30,8 +30,7 @@ struct Sm70MarlinGemmTraits { "SM70 Marlin supports CTA_N in {64, 128, 256}."); static_assert(CtaK == 16 || CtaK == 32 || CtaK == 64 || CtaK == 128, "SM70 Marlin supports CTA_K in {16, 32, 64, 128}."); - static_assert(Warps == 4 || Warps == 8, - "SM70 Marlin supports 4 or 8 warps."); + static_assert(Warps == 4 || Warps == 8, "SM70 Marlin supports 4 or 8 warps."); static_assert(PackedMacroN == 64 || PackedMacroN == 128 || PackedMacroN == 256, "SM70 Marlin packed macro-N must be 64, 128, or 256."); @@ -45,8 +44,8 @@ struct Sm70MarlinGemmTraits { using LayoutB = cutlass::layout::RowMajor; using LayoutC = cutlass::layout::RowMajor; using ThreadblockShape = cutlass::gemm::GemmShape; - using WarpShape = - typename Sm70WarpShape::Type; + using WarpShape = typename Sm70WarpShape::Type; static_assert(WarpShape::kM <= 64 && WarpShape::kN <= 64, "SM70 Marlin keeps per-warp M/N no larger than 64."); using InstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; @@ -60,9 +59,10 @@ struct Sm70MarlinGemmTraits { cutlass::MatrixShape, ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, 128 / cutlass::sizeof_bits::value>; - using IteratorB = typename Spec::template IteratorB< - ThreadblockShape, typename MmaCore::IteratorThreadMapB, GroupSize, - PackedMacroN>; + using IteratorB = + typename Spec::template IteratorB; using Mma = Sm70MarlinMmaPipelined< ThreadblockShape, IteratorA, typename MmaCore::SmemIteratorA, IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, LayoutC, @@ -102,17 +102,17 @@ torch::Tensor dispatch_sm70_marlin_group_size(Launcher const& launcher, char const* quant_name) { switch (group_size) { case -1: - return launcher.template operator()(); + return launcher.template operator()(); case 32: - return launcher.template operator()(); + return launcher.template operator()(); case 64: - return launcher.template operator()(); + return launcher.template operator()(); case 128: - return launcher.template operator()(); + return launcher.template operator()(); default: TORCH_CHECK(false, "SM70 Marlin ", quant_name, " supports only group_size -1, 32, 64, or 128. Got ", @@ -127,17 +127,17 @@ torch::Tensor dispatch_sm70_marlin_group_size_with_group32( Launcher const& launcher, int64_t group_size, char const* quant_name) { switch (group_size) { case -1: - return launcher.template operator()(); + return launcher.template operator()(); case 32: - return launcher.template operator()(); + return launcher.template operator()(); case 64: - return launcher.template operator()(); + return launcher.template operator()(); case 128: - return launcher.template operator()(); + return launcher.template operator()(); default: TORCH_CHECK(false, "SM70 Marlin ", quant_name, " supports only group_size -1, 32, 64, or 128. Got ", @@ -152,15 +152,16 @@ torch::Tensor dispatch_sm70_marlin_fp8_group_size(Launcher const& launcher, int64_t group_size) { switch (group_size) { case -1: - return launcher.template operator()(); + return launcher.template operator()(); case 128: - return launcher.template operator()(); + return launcher.template operator()(); default: - TORCH_CHECK(false, - "SM70 Marlin fp8_e4m3 supports only group_size -1 or 128. Got ", - group_size, "."); + TORCH_CHECK( + false, + "SM70 Marlin fp8_e4m3 supports only group_size -1 or 128. Got ", + group_size, "."); } return torch::Tensor(); } @@ -175,58 +176,57 @@ torch::Tensor dispatch_sm70_marlin_cta_geometry(Launcher const& launcher, return launcher.template operator()(); \ } -#define DISPATCH_SM70_GEOMETRY(CM, CN, CK, W, WM, WN, WK, PMN) \ - if (geometry.cta_m == CM && geometry.cta_n == CN && \ - geometry.cta_k == CK && geometry.warps == W && \ - geometry.warp_m == WM && geometry.warp_n == WN && \ - geometry.warp_k == WK) { \ - DISPATCH_SM70_PACKED_MACRO_N(CM, CN, CK, W, WM, WN, WK, PMN) \ +#define DISPATCH_SM70_GEOMETRY(CM, CN, CK, W, WM, WN, WK, PMN) \ + if (geometry.cta_m == CM && geometry.cta_n == CN && geometry.cta_k == CK && \ + geometry.warps == W && geometry.warp_m == WM && geometry.warp_n == WN && \ + geometry.warp_k == WK) { \ + DISPATCH_SM70_PACKED_MACRO_N(CM, CN, CK, W, WM, WN, WK, PMN) \ } -#define FOR_EACH_SM70_GEOMETRY(M) \ - M(32, 64, 32, 4, 32, 32, 16, 128) \ - M(32, 64, 32, 4, 32, 32, 16, 256) \ - M(32, 64, 64, 4, 32, 32, 32, 128) \ - M(32, 64, 64, 4, 32, 32, 32, 256) \ - M(32, 64, 64, 4, 32, 64, 16, 128) \ - M(32, 64, 64, 4, 32, 64, 16, 256) \ - M(32, 64, 128, 4, 32, 64, 32, 128) \ - M(32, 64, 128, 4, 32, 64, 32, 256) \ - M(32, 128, 32, 4, 32, 32, 32, 128) \ - M(32, 128, 32, 4, 32, 32, 32, 256) \ - M(32, 128, 64, 4, 32, 64, 32, 128) \ - M(32, 128, 64, 4, 32, 64, 32, 256) \ - M(32, 128, 64, 8, 32, 32, 32, 128) \ - M(32, 128, 64, 8, 32, 32, 32, 256) \ - M(32, 128, 128, 8, 32, 64, 32, 128) \ - M(32, 128, 128, 8, 32, 64, 32, 256) \ - M(32, 256, 32, 4, 32, 64, 32, 256) \ - M(32, 256, 64, 8, 32, 64, 32, 256) \ - M(64, 64, 32, 4, 32, 32, 32, 64) \ - M(64, 64, 32, 4, 32, 32, 32, 128) \ - M(64, 64, 32, 4, 32, 32, 32, 256) \ - M(64, 64, 32, 4, 64, 32, 16, 128) \ - M(64, 64, 32, 4, 64, 32, 16, 256) \ - M(64, 64, 64, 4, 32, 64, 32, 128) \ - M(64, 64, 64, 4, 32, 64, 32, 256) \ - M(64, 64, 64, 4, 64, 32, 32, 128) \ - M(64, 64, 64, 4, 64, 32, 32, 256) \ - M(64, 64, 64, 4, 64, 64, 16, 128) \ - M(64, 64, 64, 4, 64, 64, 16, 256) \ - M(64, 128, 32, 4, 32, 64, 32, 128) \ - M(64, 128, 32, 4, 32, 64, 32, 256) \ - M(64, 128, 32, 4, 64, 32, 32, 128) \ - M(64, 128, 32, 4, 64, 32, 32, 256) \ - M(64, 128, 128, 8, 64, 64, 32, 128) \ - M(64, 128, 128, 8, 64, 64, 32, 256) \ - M(128, 64, 32, 4, 32, 64, 32, 128) \ - M(128, 64, 32, 4, 32, 64, 32, 256) \ - M(128, 64, 128, 8, 64, 64, 32, 128) \ - M(128, 64, 128, 8, 64, 64, 32, 256) \ - M(128, 128, 32, 4, 64, 64, 32, 128) \ - M(128, 128, 32, 4, 64, 64, 32, 256) \ - M(128, 128, 64, 8, 64, 64, 32, 128) \ - M(128, 128, 64, 8, 64, 64, 32, 256) \ +#define FOR_EACH_SM70_GEOMETRY(M) \ + M(32, 64, 32, 4, 32, 32, 16, 128) \ + M(32, 64, 32, 4, 32, 32, 16, 256) \ + M(32, 64, 64, 4, 32, 32, 32, 128) \ + M(32, 64, 64, 4, 32, 32, 32, 256) \ + M(32, 64, 64, 4, 32, 64, 16, 128) \ + M(32, 64, 64, 4, 32, 64, 16, 256) \ + M(32, 64, 128, 4, 32, 64, 32, 128) \ + M(32, 64, 128, 4, 32, 64, 32, 256) \ + M(32, 128, 32, 4, 32, 32, 32, 128) \ + M(32, 128, 32, 4, 32, 32, 32, 256) \ + M(32, 128, 64, 4, 32, 64, 32, 128) \ + M(32, 128, 64, 4, 32, 64, 32, 256) \ + M(32, 128, 64, 8, 32, 32, 32, 128) \ + M(32, 128, 64, 8, 32, 32, 32, 256) \ + M(32, 128, 128, 8, 32, 64, 32, 128) \ + M(32, 128, 128, 8, 32, 64, 32, 256) \ + M(32, 256, 32, 4, 32, 64, 32, 256) \ + M(32, 256, 64, 8, 32, 64, 32, 256) \ + M(64, 64, 32, 4, 32, 32, 32, 64) \ + M(64, 64, 32, 4, 32, 32, 32, 128) \ + M(64, 64, 32, 4, 32, 32, 32, 256) \ + M(64, 64, 32, 4, 64, 32, 16, 128) \ + M(64, 64, 32, 4, 64, 32, 16, 256) \ + M(64, 64, 64, 4, 32, 64, 32, 128) \ + M(64, 64, 64, 4, 32, 64, 32, 256) \ + M(64, 64, 64, 4, 64, 32, 32, 128) \ + M(64, 64, 64, 4, 64, 32, 32, 256) \ + M(64, 64, 64, 4, 64, 64, 16, 128) \ + M(64, 64, 64, 4, 64, 64, 16, 256) \ + M(64, 128, 32, 4, 32, 64, 32, 128) \ + M(64, 128, 32, 4, 32, 64, 32, 256) \ + M(64, 128, 32, 4, 64, 32, 32, 128) \ + M(64, 128, 32, 4, 64, 32, 32, 256) \ + M(64, 128, 128, 8, 64, 64, 32, 128) \ + M(64, 128, 128, 8, 64, 64, 32, 256) \ + M(128, 64, 32, 4, 32, 64, 32, 128) \ + M(128, 64, 32, 4, 32, 64, 32, 256) \ + M(128, 64, 128, 8, 64, 64, 32, 128) \ + M(128, 64, 128, 8, 64, 64, 32, 256) \ + M(128, 128, 32, 4, 64, 64, 32, 128) \ + M(128, 128, 32, 4, 64, 64, 32, 256) \ + M(128, 128, 64, 8, 64, 64, 32, 128) \ + M(128, 128, 64, 8, 64, 64, 32, 256) \ M(128, 256, 32, 8, 64, 64, 32, 256) FOR_EACH_SM70_GEOMETRY(DISPATCH_SM70_GEOMETRY) @@ -261,8 +261,8 @@ torch::Tensor dispatch_sm70_marlin_geometry(Launcher const& launcher, int64_t group_size, char const* quant_name) { return dispatch_sm70_marlin_cta_geometry( - Sm70MarlinGroupSizeDispatchLauncher{ - launcher, group_size, quant_name}, + Sm70MarlinGroupSizeDispatchLauncher{launcher, group_size, + quant_name}, geometry, packed_macro_n, quant_name); } @@ -300,8 +300,8 @@ torch::Tensor dispatch_sm70_marlin_geometry_with_group32( Launcher const& launcher, Sm70CtaGeometry geometry, int packed_macro_n, int64_t group_size, char const* quant_name) { return dispatch_sm70_marlin_cta_geometry( - Sm70MarlinGroupSize32DispatchLauncher{ - launcher, group_size, quant_name}, + Sm70MarlinGroupSize32DispatchLauncher{launcher, group_size, + quant_name}, geometry, packed_macro_n, quant_name); } @@ -339,8 +339,8 @@ struct Sm70MarlinFixedGroupDispatchLauncher { int WarpK, int PackedMacroN> torch::Tensor operator()() const { return dispatch_sm70_marlin_fixed_group_size< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN>(inner, group_size, quant_name); + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>( + inner, group_size, quant_name); } }; diff --git a/csrc/quantization/marlin/sm70_marlin_iterator_utils.cuh b/csrc/quantization/marlin/sm70_marlin_iterator_utils.cuh index 1e134c6c68..e5ff7e3782 100644 --- a/csrc/quantization/marlin/sm70_marlin_iterator_utils.cuh +++ b/csrc/quantization/marlin/sm70_marlin_iterator_utils.cuh @@ -30,23 +30,23 @@ struct QwordVector<4> { template CUTLASS_DEVICE typename QwordVector::Type load_qword_vector( uint32_t const* ptr) { - static_assert(ContiguousWords == 1 || ContiguousWords == 2 || - ContiguousWords == 4, - "SM70 qword vector load supports 1, 2, or 4 words."); + static_assert( + ContiguousWords == 1 || ContiguousWords == 2 || ContiguousWords == 4, + "SM70 qword vector load supports 1, 2, or 4 words."); if constexpr (ContiguousWords == 1) { return *ptr; } else { - return *reinterpret_cast::Type const*>( - ptr); + return *reinterpret_cast< + typename QwordVector::Type const*>(ptr); } } template CUTLASS_DEVICE uint32_t qword_from_vector( typename QwordVector::Type const& words, int c) { - static_assert(ContiguousWords == 1 || ContiguousWords == 2 || - ContiguousWords == 4, - "SM70 qword vector access supports 1, 2, or 4 words."); + static_assert( + ContiguousWords == 1 || ContiguousWords == 2 || ContiguousWords == 4, + "SM70 qword vector access supports 1, 2, or 4 words."); if constexpr (ContiguousWords == 1) { return words; } else { @@ -68,9 +68,9 @@ CUTLASS_DEVICE uint32_t qword_from_vector(uint2 const& words, int c) { template CUTLASS_DEVICE int u4_packed_macro_n_qweight_offset_from_logical( int size_n, int logical_k, int logical_n) { - static_assert(PackedMacroN == 64 || PackedMacroN == 128 || - PackedMacroN == 256, - "SM70 Marlin packed macro-N must be 64, 128, or 256."); + static_assert( + PackedMacroN == 64 || PackedMacroN == 128 || PackedMacroN == 256, + "SM70 Marlin packed macro-N must be 64, 128, or 256."); constexpr int kGroupTiles = PackedMacroN / kQuantTileN; int const k_tile = logical_k / kQuantTileK; int const local_k = logical_k - k_tile * kQuantTileK; @@ -89,9 +89,9 @@ CUTLASS_DEVICE int u4_packed_macro_n_qweight_offset_from_logical( template CUTLASS_DEVICE int u8_packed_macro_n_qweight_offset_from_logical( int size_n, int logical_k, int logical_n) { - static_assert(PackedMacroN == 64 || PackedMacroN == 128 || - PackedMacroN == 256, - "SM70 Marlin packed macro-N must be 64, 128, or 256."); + static_assert( + PackedMacroN == 64 || PackedMacroN == 128 || PackedMacroN == 256, + "SM70 Marlin packed macro-N must be 64, 128, or 256."); constexpr int kGroupTiles = PackedMacroN / kQuantTileN; int const k_tile = logical_k / kQuantTileK; int const local_k = logical_k - k_tile * kQuantTileK; @@ -109,9 +109,9 @@ CUTLASS_DEVICE int u8_packed_macro_n_qweight_offset_from_logical( template CUTLASS_DEVICE int u8_packed_macro_n_qweight_word_stride() { - static_assert(PackedMacroN == 64 || PackedMacroN == 128 || - PackedMacroN == 256, - "SM70 Marlin packed macro-N must be 64, 128, or 256."); + static_assert( + PackedMacroN == 64 || PackedMacroN == 128 || PackedMacroN == 256, + "SM70 Marlin packed macro-N must be 64, 128, or 256."); return PackedMacroN / kQuantTileN; } diff --git a/csrc/quantization/marlin/sm70_marlin_mma.cuh b/csrc/quantization/marlin/sm70_marlin_mma.cuh index 7d3d4421bb..ded9b11707 100644 --- a/csrc/quantization/marlin/sm70_marlin_mma.cuh +++ b/csrc/quantization/marlin/sm70_marlin_mma.cuh @@ -7,17 +7,16 @@ namespace marlin::sm70 { -template < - typename Shape_, typename IteratorA_, typename SmemIteratorA_, - typename IteratorB_, typename SmemIteratorB_, typename ElementC_, - typename LayoutC_, typename Policy_, - typename TransformA_ = cutlass::NumericArrayConverter< - typename SmemIteratorA_::Element, typename IteratorA_::Element, - IteratorA_::Fragment::kElements>, - typename TransformB_ = cutlass::NumericArrayConverter< - typename SmemIteratorB_::Element, typename IteratorB_::Element, - IteratorB_::Fragment::kElements>, - typename Enable = bool> +template , + typename TransformB_ = cutlass::NumericArrayConverter< + typename SmemIteratorB_::Element, typename IteratorB_::Element, + IteratorB_::Fragment::kElements>, + typename Enable = bool> class Sm70MarlinMmaPipelined : public cutlass::gemm::threadblock::MmaBase { public: @@ -106,10 +105,10 @@ class Sm70MarlinMmaPipelined public: CUTLASS_DEVICE - Sm70MarlinMmaPipelined( - typename Base::SharedStorage& shared_storage, int thread_idx, - int warp_idx, int lane_idx, TransformA transform_A = TransformA(), - TransformB transform_B = TransformB()) + Sm70MarlinMmaPipelined(typename Base::SharedStorage& shared_storage, + int thread_idx, int warp_idx, int lane_idx, + TransformA transform_A = TransformA(), + TransformB transform_B = TransformB()) : Base(shared_storage, thread_idx, warp_idx, lane_idx), smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx), @@ -154,12 +153,12 @@ class Sm70MarlinMmaPipelined this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); if constexpr (Policy::kPartitionsK > 1) { - add_warp_tile_k_group_offset( - (Policy::kPartitionsK - 1) * Base::kWarpGemmIterations); + add_warp_tile_k_group_offset((Policy::kPartitionsK - 1) * + Base::kWarpGemmIterations); } } else { - add_warp_tile_k_group_offset( - -(Policy::kPartitionsK + 1) * Base::kWarpGemmIterations); + add_warp_tile_k_group_offset(-(Policy::kPartitionsK + 1) * + Base::kWarpGemmIterations); } smem_write_stage_idx ^= 1; @@ -222,10 +221,10 @@ class Sm70MarlinMmaPipelined advance_smem_stages(); } - this->warp_tile_iterator_A_.set_kgroup_index( - (warp_mma_k + 1) % Base::kWarpGemmIterations); - this->warp_tile_iterator_B_.set_kgroup_index( - (warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % + Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % + Base::kWarpGemmIterations); this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); diff --git a/csrc/quantization/marlin/sm70_marlin_mxfp4_gemm.cu b/csrc/quantization/marlin/sm70_marlin_mxfp4_gemm.cu index 8a22142725..8235c2b17c 100644 --- a/csrc/quantization/marlin/sm70_marlin_mxfp4_gemm.cu +++ b/csrc/quantization/marlin/sm70_marlin_mxfp4_gemm.cu @@ -21,25 +21,25 @@ #include "cutlass/layout/tensor_op_multiplicand_sm70.h" #include "cutlass/transform/threadblock/predicated_tile_iterator.h" -using marlin::sm70::Sm70CtaGeometry; -using marlin::sm70::Sm70AtomicFp16Epilogue; -using marlin::sm70::Sm70AtomicFp32Epilogue; -using marlin::sm70::sm70_marlin_narrow_f32_to_f16; -using marlin::sm70::Sm70MarlinGemmTraits; -using marlin::sm70::Sm70SplitKPartition; -using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; -using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; using marlin::sm70::configure_sm70_dynamic_smem; using marlin::sm70::dispatch_sm70_marlin_fixed_group_geometry; using marlin::sm70::kQuantTileK; using marlin::sm70::kQuantTileN; -using marlin::sm70::sm70_marlin_dense_auto_params; using marlin::sm70::load_qword_vector; using marlin::sm70::qword_from_vector; -using marlin::sm70::sm70_marlin_cta_grid; using marlin::sm70::sm70_active_split_k; +using marlin::sm70::sm70_marlin_cta_grid; +using marlin::sm70::sm70_marlin_dense_auto_params; +using marlin::sm70::sm70_marlin_narrow_f32_to_f16; using marlin::sm70::sm70_splitk_partition; +using marlin::sm70::Sm70AtomicFp16Epilogue; +using marlin::sm70::Sm70AtomicFp32Epilogue; +using marlin::sm70::Sm70CtaGeometry; +using marlin::sm70::Sm70MarlinGemmTraits; +using marlin::sm70::Sm70SplitKPartition; using marlin::sm70::u4_packed_macro_n_qweight_offset_from_logical; +using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; +using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; namespace { @@ -72,7 +72,8 @@ void dequant_e8m0_scales_to_half2(int q, half2* scale_cache) { *reinterpret_cast(scale_cache) = vec_res; } -template +template class Sm70Mxfp4IteratorB { public: using Shape = Shape_; @@ -80,12 +81,11 @@ class Sm70Mxfp4IteratorB { static int const kGroupSize = GroupSize_; static int const kPackedMacroN = PackedMacroN_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + using Fragment = cutlass::Array; static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, "SM70 Marlin MXFP4 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, + static_assert(ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, "SM70 Marlin MXFP4 IteratorB expects one contiguous iteration " "per 64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, @@ -94,7 +94,8 @@ class Sm70Mxfp4IteratorB { "SM70 Marlin MXFP4 IteratorB expects one packed FP4 word per " "access."); static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin U4-family IteratorB expects one or more strided iterations."); + "SM70 Marlin U4-family IteratorB expects one or more strided " + "iterations."); struct Params { int size_n; int aligned_initial_k_step; @@ -176,8 +177,8 @@ class Sm70Mxfp4IteratorB { CUTLASS_DEVICE static int qweight_offset_from_logical(Params const& params, int logical_k, int logical_n) { - return u4_packed_macro_n_qweight_offset_from_logical(params.size_n, logical_k, - logical_n); + return u4_packed_macro_n_qweight_offset_from_logical( + params.size_n, logical_k, logical_n); } CUTLASS_DEVICE @@ -196,14 +197,12 @@ class Sm70Mxfp4IteratorB { CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { int const logical_k = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + k_offset_ + thread_offset_.strided() + s * ThreadMap::Delta::kStrided; int const group = scale_group(logical_k); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; int const cache_index = c + s * ThreadMap::Iterations::kContiguous; cache_metadata_e8m0_scales(cache_index, group, cache_n); } @@ -256,9 +255,11 @@ class Sm70Mxfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MXFP4 contiguous iteration count."); - uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_offset_); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MXFP4 contiguous iteration count."); + uint32_t const qword = + load_qword_vector<1>(qweight_ + qweight_base_offset_); half2 const* scale_vec = cached_scales_; half2 deq[2]; @@ -276,14 +277,12 @@ class Sm70Mxfp4IteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; int const qweight_base_s = qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - auto const qwords = - load_qword_vector<4>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<4>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -304,8 +303,7 @@ class Sm70Mxfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - auto const qwords = - load_qword_vector<2>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<2>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -326,9 +324,11 @@ class Sm70Mxfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin MXFP4 contiguous iteration count."); - uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_s); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin MXFP4 contiguous iteration count."); + uint32_t const qword = + load_qword_vector<1>(qweight_ + qweight_base_s); constexpr int kAccess = ThreadMap::kElementsPerAccess; int const frag_base = s * kAccess; half2 const* scale_vec = cached_scales_; @@ -361,26 +361,25 @@ class Sm70Mxfp4IteratorB { struct Sm70Mxfp4GemmSpec { template - using IteratorB = Sm70Mxfp4IteratorB; + using IteratorB = + Sm70Mxfp4IteratorB; }; -template +template using Sm70Mxfp4GemmTraits = - Sm70MarlinGemmTraits; + Sm70MarlinGemmTraits; -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_mxfp4_gemm_kernel( +template +__global__ __launch_bounds__(Warps * 32, 1) void sm70_marlin_mxfp4_gemm_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, - uint8_t const* __restrict__ b_scales, - cutlass::half_t* __restrict__ c, int m, int n, int k, int lda) { - using Traits = - Sm70Mxfp4GemmTraits; + uint8_t const* __restrict__ b_scales, cutlass::half_t* __restrict__ c, + int m, int n, int k, int lda) { + using Traits = Sm70Mxfp4GemmTraits; using Mma = typename Traits::Mma; using Epilogue = typename Traits::Epilogue; @@ -400,10 +399,10 @@ void sm70_marlin_mxfp4_gemm_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -428,16 +427,17 @@ void sm70_marlin_mxfp4_gemm_kernel( epilogue(output_op, iterator_D, accumulators, iterator_C); } -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_mxfp4_gemm_splitk_kernel( +template +__global__ +__launch_bounds__(Warps * 32, 1) void sm70_marlin_mxfp4_gemm_splitk_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, - uint8_t const* __restrict__ b_scales, - cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, int k, int lda, int requested_split_k) { - using Traits = Sm70Mxfp4GemmTraits; + uint8_t const* __restrict__ b_scales, cutlass::half_t* __restrict__ c, + float* __restrict__ c32, int m, int n, int k, int lda, + int requested_split_k) { + using Traits = Sm70Mxfp4GemmTraits; using Mma = typename Traits::Mma; using AtomicEpilogue = Sm70AtomicFp32Epilogue; @@ -448,14 +448,13 @@ void sm70_marlin_mxfp4_gemm_splitk_kernel( int const thread_idx = threadIdx.x; int const warp_idx = cutlass::canonical_warp_idx_sync(); int const lane_idx = threadIdx.x % 32; - Sm70SplitKPartition const partition = - sm70_splitk_partition(k, requested_split_k, int(blockIdx.z)); + Sm70SplitKPartition const partition = sm70_splitk_partition( + k, requested_split_k, int(blockIdx.z)); if (partition.partition_k == 0) { return; } - cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, - partition.k_begin}; + cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, partition.k_begin}; cutlass::MatrixCoord tb_offset_B{partition.k_begin, int(blockIdx.y) * CtaN}; cutlass::MatrixCoord tb_offset_C{int(blockIdx.x) * CtaM, int(blockIdx.y) * CtaN}; @@ -463,10 +462,10 @@ void sm70_marlin_mxfp4_gemm_splitk_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -480,9 +479,8 @@ void sm70_marlin_mxfp4_gemm_splitk_kernel( mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); typename AtomicEpilogue::OutputTileIterator iterator_D( - typename AtomicEpilogue::OutputTileIterator::Params(layout_c), - c, cutlass::MatrixCoord(m, n), - thread_idx, tb_offset_C); + typename AtomicEpilogue::OutputTileIterator::Params(layout_c), c, + cutlass::MatrixCoord(m, n), thread_idx, tb_offset_C); AtomicEpilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx); @@ -491,18 +489,20 @@ void sm70_marlin_mxfp4_gemm_splitk_kernel( } // namespace -template -torch::Tensor launch_sm70_marlin_mxfp4_gemm( - torch::Tensor& a, torch::Tensor& c, torch::Tensor& b_q_weight, - torch::Tensor& b_scales, int64_t size_m, int64_t size_n, int64_t size_k, - int requested_split_k) { +template +torch::Tensor launch_sm70_marlin_mxfp4_gemm(torch::Tensor& a, torch::Tensor& c, + torch::Tensor& b_q_weight, + torch::Tensor& b_scales, + int64_t size_m, int64_t size_n, + int64_t size_k, + int requested_split_k) { auto kernel = - sm70_marlin_mxfp4_gemm_kernel; - using SharedStorage = typename Sm70Mxfp4GemmTraits< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN>::SharedStorage; + sm70_marlin_mxfp4_gemm_kernel; + using SharedStorage = + typename Sm70Mxfp4GemmTraits::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -522,21 +522,19 @@ torch::Tensor launch_sm70_marlin_mxfp4_gemm( } TORCH_CHECK(size_k % int64_t(CtaK) == 0, - "SM70 Marlin mxfp4 requires K divisible by CTA_K=", - CtaK, " for requested_split_k > 1. Got K=", size_k, + "SM70 Marlin mxfp4 requires K divisible by CTA_K=", CtaK, + " for requested_split_k > 1. Got K=", size_k, ", requested_split_k=", requested_split_k, "."); - auto split_kernel = - sm70_marlin_mxfp4_gemm_splitk_kernel; + auto split_kernel = sm70_marlin_mxfp4_gemm_splitk_kernel< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; smem_bytes = configure_sm70_dynamic_smem(split_kernel); int64_t const numel = size_m * size_n; // 0009: fp32 split-K scratch (zeroed for the fp32 atomicAdd), M*N floats. // Only allocated on the split_k>1 path, which the auto-selector uses ONLY at // low batch (M<128), so the scratch is bounded and small. - auto c32 = torch::zeros({size_m, size_n}, - a.options().dtype(at::kFloat)); + auto c32 = torch::zeros({size_m, size_n}, a.options().dtype(at::kFloat)); dim3 grid = sm70_marlin_cta_grid(size_m, size_n, CtaM, CtaN); int const active_split_k = @@ -547,9 +545,9 @@ torch::Tensor launch_sm70_marlin_mxfp4_gemm( reinterpret_cast(b_q_weight.data_ptr()), reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(c.data_ptr()), - c32.data_ptr(), - static_cast(size_m), static_cast(size_n), - static_cast(size_k), static_cast(a.stride(0)), requested_split_k); + c32.data_ptr(), static_cast(size_m), static_cast(size_n), + static_cast(size_k), static_cast(a.stride(0)), + requested_split_k); C10_CUDA_KERNEL_LAUNCH_CHECK(); // 0009: single deterministic-width fp32->fp16 narrow of the scratch. @@ -557,8 +555,8 @@ torch::Tensor launch_sm70_marlin_mxfp4_gemm( int const narrow_grid = static_cast((numel + narrow_block - 1) / narrow_block); sm70_marlin_narrow_f32_to_f16<256><<>>( - c32.data_ptr(), - reinterpret_cast(c.data_ptr()), numel); + c32.data_ptr(), reinterpret_cast(c.data_ptr()), + numel); C10_CUDA_KERNEL_LAUNCH_CHECK(); return c; @@ -574,11 +572,11 @@ struct Sm70Mxfp4Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - return launch_sm70_marlin_mxfp4_gemm( + return launch_sm70_marlin_mxfp4_gemm( a, c, b_q_weight, b_scales, size_m, size_n, size_k, requested_split_k); } }; @@ -590,14 +588,16 @@ torch::Tensor sm70_marlin_mxfp4_gemm(torch::Tensor& a, torch::Tensor& c, int64_t group_size) { c10::cuda::CUDAGuard device_guard(a.device()); - auto const params = sm70_marlin_dense_auto_params( - "mxfp4", group_size, size_m, size_n, size_k); + auto const params = sm70_marlin_dense_auto_params("mxfp4", group_size, size_m, + size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_dense_cta_geometry_supported("SM70 Marlin mxfp4", geometry); - validate_sm70_marlin_dense_cta_n_alignment("SM70 Marlin mxfp4", geometry, size_n); + validate_sm70_marlin_dense_cta_geometry_supported("SM70 Marlin mxfp4", + geometry); + validate_sm70_marlin_dense_cta_n_alignment("SM70 Marlin mxfp4", geometry, + size_n); Sm70Mxfp4Launcher const launcher{ - a, c, b_q_weight, b_scales, size_m, size_n, size_k, - params.requested_split_k}; + a, c, b_q_weight, b_scales, + size_m, size_n, size_k, params.requested_split_k}; return dispatch_sm70_marlin_fixed_group_geometry<32>( launcher, geometry, params.packed_macro_n, group_size, "mxfp4"); } diff --git a/csrc/quantization/marlin/sm70_marlin_nvfp4_gemm.cu b/csrc/quantization/marlin/sm70_marlin_nvfp4_gemm.cu index bd33da2e22..debb47d16c 100644 --- a/csrc/quantization/marlin/sm70_marlin_nvfp4_gemm.cu +++ b/csrc/quantization/marlin/sm70_marlin_nvfp4_gemm.cu @@ -21,31 +21,32 @@ #include "cutlass/layout/tensor_op_multiplicand_sm70.h" #include "cutlass/transform/threadblock/predicated_tile_iterator.h" -using marlin::sm70::Sm70CtaGeometry; -using marlin::sm70::Sm70AtomicFp16Epilogue; -using marlin::sm70::Sm70AtomicFp32Epilogue; -using marlin::sm70::sm70_marlin_narrow_f32_to_f16; -using marlin::sm70::Sm70MarlinGemmTraits; -using marlin::sm70::Sm70SplitKPartition; -using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; -using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; using marlin::sm70::configure_sm70_dynamic_smem; using marlin::sm70::dispatch_sm70_marlin_fixed_group_geometry; using marlin::sm70::kQuantTileK; using marlin::sm70::kQuantTileN; -using marlin::sm70::sm70_marlin_dense_auto_params; using marlin::sm70::load_qword_vector; using marlin::sm70::qword_from_vector; -using marlin::sm70::sm70_marlin_cta_grid; using marlin::sm70::sm70_active_split_k; +using marlin::sm70::sm70_marlin_cta_grid; +using marlin::sm70::sm70_marlin_dense_auto_params; +using marlin::sm70::sm70_marlin_narrow_f32_to_f16; using marlin::sm70::sm70_splitk_partition; +using marlin::sm70::Sm70AtomicFp16Epilogue; +using marlin::sm70::Sm70AtomicFp32Epilogue; +using marlin::sm70::Sm70CtaGeometry; +using marlin::sm70::Sm70MarlinGemmTraits; +using marlin::sm70::Sm70SplitKPartition; using marlin::sm70::u4_packed_macro_n_qweight_offset_from_logical; +using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; +using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; namespace { constexpr int kNvfp4ValuesPerWord = 8; -template +template class Sm70Nvfp4IteratorB { public: using Shape = Shape_; @@ -53,12 +54,11 @@ class Sm70Nvfp4IteratorB { static int const kGroupSize = GroupSize_; static int const kPackedMacroN = PackedMacroN_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + using Fragment = cutlass::Array; static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, "SM70 Marlin NVFP4 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, + static_assert(ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, "SM70 Marlin NVFP4 IteratorB expects one contiguous iteration " "per 64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, @@ -67,7 +67,8 @@ class Sm70Nvfp4IteratorB { "SM70 Marlin NVFP4 IteratorB expects one packed FP4 word per " "access."); static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin U4-family IteratorB expects one or more strided iterations."); + "SM70 Marlin U4-family IteratorB expects one or more strided " + "iterations."); struct Params { int size_n; int aligned_initial_k_step; @@ -149,8 +150,8 @@ class Sm70Nvfp4IteratorB { CUTLASS_DEVICE static int qweight_offset_from_logical(Params const& params, int logical_k, int logical_n) { - return u4_packed_macro_n_qweight_offset_from_logical(params.size_n, logical_k, - logical_n); + return u4_packed_macro_n_qweight_offset_from_logical( + params.size_n, logical_k, logical_n); } CUTLASS_DEVICE @@ -170,14 +171,12 @@ class Sm70Nvfp4IteratorB { CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { int const logical_k = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + k_offset_ + thread_offset_.strided() + s * ThreadMap::Delta::kStrided; int const group = scale_group(logical_k); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; int const cache_index = c + s * ThreadMap::Iterations::kContiguous; cache_metadata_fp8_scales(cache_index, group, cache_n); } @@ -230,9 +229,11 @@ class Sm70Nvfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin NVFP4 contiguous iteration count."); - uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_offset_); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin NVFP4 contiguous iteration count."); + uint32_t const qword = + load_qword_vector<1>(qweight_ + qweight_base_offset_); half2 const* scale_vec = cached_scales_; half2 deq[2]; @@ -250,14 +251,12 @@ class Sm70Nvfp4IteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; int const qweight_base_s = qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - auto const qwords = - load_qword_vector<4>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<4>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -280,8 +279,7 @@ class Sm70Nvfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - auto const qwords = - load_qword_vector<2>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<2>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -304,9 +302,11 @@ class Sm70Nvfp4IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin NVFP4 contiguous iteration count."); - uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_s); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin NVFP4 contiguous iteration count."); + uint32_t const qword = + load_qword_vector<1>(qweight_ + qweight_base_s); constexpr int kAccess = ThreadMap::kElementsPerAccess; int const frag_base = s * kAccess; half2 const* scale_vec = @@ -340,27 +340,26 @@ class Sm70Nvfp4IteratorB { struct Sm70Nvfp4GemmSpec { template - using IteratorB = Sm70Nvfp4IteratorB; + using IteratorB = + Sm70Nvfp4IteratorB; }; -template +template using Sm70Nvfp4GemmTraits = - Sm70MarlinGemmTraits; + Sm70MarlinGemmTraits; -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_nvfp4_gemm_kernel( +template +__global__ __launch_bounds__(Warps * 32, 1) void sm70_marlin_nvfp4_gemm_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, uint8_t const* __restrict__ b_scales, - float const* __restrict__ global_scale, - cutlass::half_t* __restrict__ c, int m, int n, int k, int lda) { - using Traits = - Sm70Nvfp4GemmTraits; + float const* __restrict__ global_scale, cutlass::half_t* __restrict__ c, + int m, int n, int k, int lda) { + using Traits = Sm70Nvfp4GemmTraits; using Mma = typename Traits::Mma; using Epilogue = typename Traits::Epilogue; @@ -380,10 +379,10 @@ void sm70_marlin_nvfp4_gemm_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -408,17 +407,18 @@ void sm70_marlin_nvfp4_gemm_kernel( epilogue(output_op, iterator_D, accumulators, iterator_C); } -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_nvfp4_gemm_splitk_kernel( +template +__global__ +__launch_bounds__(Warps * 32, 1) void sm70_marlin_nvfp4_gemm_splitk_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, uint8_t const* __restrict__ b_scales, - float const* __restrict__ global_scale, - cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, int k, int lda, int requested_split_k) { - using Traits = Sm70Nvfp4GemmTraits; + float const* __restrict__ global_scale, cutlass::half_t* __restrict__ c, + float* __restrict__ c32, int m, int n, int k, int lda, + int requested_split_k) { + using Traits = Sm70Nvfp4GemmTraits; using Mma = typename Traits::Mma; using AtomicEpilogue = Sm70AtomicFp32Epilogue; @@ -429,14 +429,13 @@ void sm70_marlin_nvfp4_gemm_splitk_kernel( int const thread_idx = threadIdx.x; int const warp_idx = cutlass::canonical_warp_idx_sync(); int const lane_idx = threadIdx.x % 32; - Sm70SplitKPartition const partition = - sm70_splitk_partition(k, requested_split_k, int(blockIdx.z)); + Sm70SplitKPartition const partition = sm70_splitk_partition( + k, requested_split_k, int(blockIdx.z)); if (partition.partition_k == 0) { return; } - cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, - partition.k_begin}; + cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, partition.k_begin}; cutlass::MatrixCoord tb_offset_B{partition.k_begin, int(blockIdx.y) * CtaN}; cutlass::MatrixCoord tb_offset_C{int(blockIdx.x) * CtaM, int(blockIdx.y) * CtaN}; @@ -444,10 +443,10 @@ void sm70_marlin_nvfp4_gemm_splitk_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -463,9 +462,8 @@ void sm70_marlin_nvfp4_gemm_splitk_kernel( accumulators = scale_accumulators(accumulators, global_scale[0]); typename AtomicEpilogue::OutputTileIterator iterator_D( - typename AtomicEpilogue::OutputTileIterator::Params(layout_c), - c, cutlass::MatrixCoord(m, n), - thread_idx, tb_offset_C); + typename AtomicEpilogue::OutputTileIterator::Params(layout_c), c, + cutlass::MatrixCoord(m, n), thread_idx, tb_offset_C); AtomicEpilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx); @@ -474,18 +472,18 @@ void sm70_marlin_nvfp4_gemm_splitk_kernel( } // namespace -template +template torch::Tensor launch_sm70_marlin_nvfp4_gemm( torch::Tensor& a, torch::Tensor& c, torch::Tensor& b_q_weight, torch::Tensor& b_scales, torch::Tensor& global_scale, int64_t size_m, int64_t size_n, int64_t size_k, int requested_split_k) { auto kernel = - sm70_marlin_nvfp4_gemm_kernel; - using SharedStorage = typename Sm70Nvfp4GemmTraits< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN>::SharedStorage; + sm70_marlin_nvfp4_gemm_kernel; + using SharedStorage = + typename Sm70Nvfp4GemmTraits::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -506,21 +504,19 @@ torch::Tensor launch_sm70_marlin_nvfp4_gemm( } TORCH_CHECK(size_k % int64_t(CtaK) == 0, - "SM70 Marlin nvfp4 requires K divisible by CTA_K=", - CtaK, " for requested_split_k > 1. Got K=", size_k, + "SM70 Marlin nvfp4 requires K divisible by CTA_K=", CtaK, + " for requested_split_k > 1. Got K=", size_k, ", requested_split_k=", requested_split_k, "."); - auto split_kernel = - sm70_marlin_nvfp4_gemm_splitk_kernel; + auto split_kernel = sm70_marlin_nvfp4_gemm_splitk_kernel< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN>; smem_bytes = configure_sm70_dynamic_smem(split_kernel); int64_t const numel = size_m * size_n; // 0009: fp32 split-K scratch (zeroed for the fp32 atomicAdd), M*N floats. // Only allocated on the split_k>1 path, which the auto-selector uses ONLY at // low batch (M<128), so the scratch is bounded and small. - auto c32 = torch::zeros({size_m, size_n}, - a.options().dtype(at::kFloat)); + auto c32 = torch::zeros({size_m, size_n}, a.options().dtype(at::kFloat)); dim3 grid = sm70_marlin_cta_grid(size_m, size_n, CtaM, CtaN); int const active_split_k = @@ -532,9 +528,9 @@ torch::Tensor launch_sm70_marlin_nvfp4_gemm( reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(global_scale.data_ptr()), reinterpret_cast(c.data_ptr()), - c32.data_ptr(), - static_cast(size_m), static_cast(size_n), - static_cast(size_k), static_cast(a.stride(0)), requested_split_k); + c32.data_ptr(), static_cast(size_m), static_cast(size_n), + static_cast(size_k), static_cast(a.stride(0)), + requested_split_k); C10_CUDA_KERNEL_LAUNCH_CHECK(); // 0009: single deterministic-width fp32->fp16 narrow of the scratch. @@ -542,8 +538,8 @@ torch::Tensor launch_sm70_marlin_nvfp4_gemm( int const narrow_grid = static_cast((numel + narrow_block - 1) / narrow_block); sm70_marlin_narrow_f32_to_f16<256><<>>( - c32.data_ptr(), - reinterpret_cast(c.data_ptr()), numel); + c32.data_ptr(), reinterpret_cast(c.data_ptr()), + numel); C10_CUDA_KERNEL_LAUNCH_CHECK(); return c; @@ -560,11 +556,11 @@ struct Sm70Nvfp4Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - return launch_sm70_marlin_nvfp4_gemm( + return launch_sm70_marlin_nvfp4_gemm( a, c, b_q_weight, b_scales, global_scale, size_m, size_n, size_k, requested_split_k); } @@ -578,14 +574,22 @@ torch::Tensor sm70_marlin_nvfp4_gemm(torch::Tensor& a, torch::Tensor& c, int64_t size_k, int64_t group_size) { c10::cuda::CUDAGuard device_guard(a.device()); - auto const params = sm70_marlin_dense_auto_params( - "nvfp4", group_size, size_m, size_n, size_k); + auto const params = sm70_marlin_dense_auto_params("nvfp4", group_size, size_m, + size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_dense_cta_geometry_supported("SM70 Marlin nvfp4", geometry); - validate_sm70_marlin_dense_cta_n_alignment("SM70 Marlin nvfp4", geometry, size_n); - Sm70Nvfp4Launcher const launcher{ - a, c, b_q_weight, b_scales, global_scale, size_m, size_n, size_k, - params.requested_split_k}; + validate_sm70_marlin_dense_cta_geometry_supported("SM70 Marlin nvfp4", + geometry); + validate_sm70_marlin_dense_cta_n_alignment("SM70 Marlin nvfp4", geometry, + size_n); + Sm70Nvfp4Launcher const launcher{a, + c, + b_q_weight, + b_scales, + global_scale, + size_m, + size_n, + size_k, + params.requested_split_k}; return dispatch_sm70_marlin_fixed_group_geometry<16>( launcher, geometry, params.packed_macro_n, group_size, "nvfp4"); } diff --git a/csrc/quantization/marlin/sm70_marlin_splitk.cuh b/csrc/quantization/marlin/sm70_marlin_splitk.cuh index a0fb01910a..463f08b84f 100644 --- a/csrc/quantization/marlin/sm70_marlin_splitk.cuh +++ b/csrc/quantization/marlin/sm70_marlin_splitk.cuh @@ -22,9 +22,7 @@ int sm70_splitk_ceil_div_int(int numerator, int denominator) { } CUTLASS_HOST_DEVICE -int sm70_splitk_min_int(int lhs, int rhs) { - return lhs < rhs ? lhs : rhs; -} +int sm70_splitk_min_int(int lhs, int rhs) { return lhs < rhs ? lhs : rhs; } template CUTLASS_HOST_DEVICE int sm70_splitk_group_tiles() { @@ -76,8 +74,8 @@ int sm70_splitk_partition_tile_count(int remaining_tiles, } template -CUTLASS_HOST_DEVICE Sm70SplitKPartition sm70_splitk_partition( - int k, int requested_split_k, int partition_idx) { +CUTLASS_HOST_DEVICE Sm70SplitKPartition +sm70_splitk_partition(int k, int requested_split_k, int partition_idx) { int const active_split_k = sm70_active_split_k(k, requested_split_k, CtaK); if (partition_idx >= active_split_k) { return {0, 0}; @@ -135,10 +133,9 @@ class Sm70AtomicFp16Epilogue { int const frag_row_idx = row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster); - int const row_offset = - row * ThreadMap::Delta::kRow + - group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; + int const row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; int const logical_row = thread_start_row + row_offset; CUTLASS_PRAGMA_UNROLL @@ -168,7 +165,7 @@ class Sm70AtomicFp16Epilogue { public: CUTLASS_DEVICE Sm70AtomicFp16Epilogue(SharedStorage& shared_storage, int thread_idx, - int warp_idx, int lane_idx) + int warp_idx, int lane_idx) : warp_tile_iterator_(shared_storage.reference(), lane_idx), shared_load_iterator_(shared_storage.reference(), thread_idx) { using WarpCount = typename CutlassEpilogue::WarpCount; @@ -177,8 +174,7 @@ class Sm70AtomicFp16Epilogue { int const warp_m = warp_mn % WarpCount::kM; int const warp_n = warp_mn / WarpCount::kM; - cutlass::MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, - warp_n}; + cutlass::MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, warp_n}; warp_tile_iterator_.add_tile_offset(warp_offset); } @@ -220,8 +216,7 @@ class Sm70AtomicFp16Epilogue { CutlassEpilogue::kSmemPointerOffset); } - atomic_store_fragment(destination_iterator, aligned_accum_fragment, c, - n); + atomic_store_fragment(destination_iterator, aligned_accum_fragment, c, n); ++destination_iterator; } } @@ -263,10 +258,9 @@ class Sm70AtomicFp32Epilogue { int const frag_row_idx = row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster); - int const row_offset = - row * ThreadMap::Delta::kRow + - group * ThreadMap::Delta::kGroup + - cluster * ThreadMap::Delta::kCluster; + int const row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; int const logical_row = thread_start_row + row_offset; CUTLASS_PRAGMA_UNROLL @@ -295,7 +289,7 @@ class Sm70AtomicFp32Epilogue { public: CUTLASS_DEVICE Sm70AtomicFp32Epilogue(SharedStorage& shared_storage, int thread_idx, - int warp_idx, int lane_idx) + int warp_idx, int lane_idx) : warp_tile_iterator_(shared_storage.reference(), lane_idx), shared_load_iterator_(shared_storage.reference(), thread_idx) { using WarpCount = typename CutlassEpilogue::WarpCount; @@ -304,15 +298,14 @@ class Sm70AtomicFp32Epilogue { int const warp_m = warp_mn % WarpCount::kM; int const warp_n = warp_mn / WarpCount::kM; - cutlass::MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, - warp_n}; + cutlass::MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, warp_n}; warp_tile_iterator_.add_tile_offset(warp_offset); } CUTLASS_DEVICE void operator()(OutputTileIterator destination_iterator, - AccumulatorTile const& accumulators, - float* __restrict__ c, int n) { + AccumulatorTile const& accumulators, float* __restrict__ c, + int n) { AccumulatorFragmentIterator accum_fragment_iterator(accumulators); CUTLASS_PRAGMA_UNROLL @@ -347,17 +340,16 @@ class Sm70AtomicFp32Epilogue { CutlassEpilogue::kSmemPointerOffset); } - atomic_store_fragment(destination_iterator, aligned_accum_fragment, c, - n); + atomic_store_fragment(destination_iterator, aligned_accum_fragment, c, n); ++destination_iterator; } } }; - -// 0009: single fp32->fp16 narrow of the split-K scratch (replaces the per-partial -// __float2half_rn narrow the fp16-atomic epilogue did S times). Templated so the -// symbol has vague linkage across the 7 gemm TUs that include this header. +// 0009: single fp32->fp16 narrow of the split-K scratch (replaces the +// per-partial +// __float2half_rn narrow the fp16-atomic epilogue did S times). Templated so +// the symbol has vague linkage across the 7 gemm TUs that include this header. template __global__ void sm70_marlin_narrow_f32_to_f16(const float* __restrict__ src, half* __restrict__ dst, diff --git a/csrc/quantization/marlin/sm70_marlin_u4_gemm.cu b/csrc/quantization/marlin/sm70_marlin_u4_gemm.cu index 0037ae8187..e2038fae64 100644 --- a/csrc/quantization/marlin/sm70_marlin_u4_gemm.cu +++ b/csrc/quantization/marlin/sm70_marlin_u4_gemm.cu @@ -20,25 +20,25 @@ #include "cutlass/layout/tensor_op_multiplicand_sm70.h" #include "cutlass/transform/threadblock/predicated_tile_iterator.h" -using marlin::sm70::Sm70CtaGeometry; -using marlin::sm70::Sm70AtomicFp16Epilogue; -using marlin::sm70::Sm70AtomicFp32Epilogue; -using marlin::sm70::sm70_marlin_narrow_f32_to_f16; -using marlin::sm70::Sm70MarlinGemmTraits; -using marlin::sm70::Sm70SplitKPartition; -using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; -using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; using marlin::sm70::configure_sm70_dynamic_smem; using marlin::sm70::dispatch_sm70_marlin_geometry; using marlin::sm70::kQuantTileK; using marlin::sm70::kQuantTileN; -using marlin::sm70::sm70_marlin_dense_auto_params; using marlin::sm70::load_qword_vector; using marlin::sm70::qword_from_vector; -using marlin::sm70::sm70_marlin_cta_grid; using marlin::sm70::sm70_active_split_k; +using marlin::sm70::sm70_marlin_cta_grid; +using marlin::sm70::sm70_marlin_dense_auto_params; +using marlin::sm70::sm70_marlin_narrow_f32_to_f16; using marlin::sm70::sm70_splitk_partition; +using marlin::sm70::Sm70AtomicFp16Epilogue; +using marlin::sm70::Sm70AtomicFp32Epilogue; +using marlin::sm70::Sm70CtaGeometry; +using marlin::sm70::Sm70MarlinGemmTraits; +using marlin::sm70::Sm70SplitKPartition; using marlin::sm70::u4_packed_macro_n_qweight_offset_from_logical; +using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; +using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; namespace { @@ -54,12 +54,11 @@ class Sm70U4ZpIteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + using Fragment = cutlass::Array; static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, "SM70 Marlin U4 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, + static_assert(ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, "SM70 Marlin U4 IteratorB expects one contiguous iteration per " "64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, @@ -68,7 +67,8 @@ class Sm70U4ZpIteratorB { "SM70 Marlin U4 IteratorB expects one packed int4 word per " "access."); static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin U4-family IteratorB expects one or more strided iterations."); + "SM70 Marlin U4-family IteratorB expects one or more strided " + "iterations."); struct Params { int size_n; int aligned_initial_k_step; @@ -115,7 +115,7 @@ class Sm70U4ZpIteratorB { qweight_base_offset_ = qweight_offset_from_logical(params_, logical_k, logical_n); if constexpr (kGroupSize == -1) { - cache_current_group_metadata(0); + cache_current_group_metadata(0); } } @@ -151,8 +151,7 @@ class Sm70U4ZpIteratorB { if constexpr (kGroupSize == -1) { return 0; } else { - static_assert(kGroupSize == 32 || kGroupSize == 64 || - kGroupSize == 128, + static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, "SM70 Marlin U4 supports only group sizes " "-1, 32, 64, and 128."); return logical_k / kGroupSize; @@ -162,8 +161,8 @@ class Sm70U4ZpIteratorB { CUTLASS_DEVICE static int qweight_offset_from_logical(Params const& params, int logical_k, int logical_n) { - return u4_packed_macro_n_qweight_offset_from_logical(params.size_n, logical_k, - logical_n); + return u4_packed_macro_n_qweight_offset_from_logical( + params.size_n, logical_k, logical_n); } CUTLASS_DEVICE @@ -176,8 +175,8 @@ class Sm70U4ZpIteratorB { scale_cache[2] = scale_vec[2]; scale_cache[3] = scale_vec[3]; - half2 const* zp_vec = reinterpret_cast( - zp_ + group * params_.size_n + cache_n); + half2 const* zp_vec = + reinterpret_cast(zp_ + group * params_.size_n + cache_n); half2* zp_cache = cached_zp_ + c * 4; zp_cache[0] = zp_vec[0]; zp_cache[1] = zp_vec[1]; @@ -196,8 +195,8 @@ class Sm70U4ZpIteratorB { scale_cache[2] = scale_vec[2]; scale_cache[3] = scale_vec[3]; - uint4 const zp_words = *reinterpret_cast( - zp_ + group * params_.size_n + cache_n); + uint4 const zp_words = + *reinterpret_cast(zp_ + group * params_.size_n + cache_n); half2 const* zp_vec = reinterpret_cast(&zp_words); half2* zp_cache = cached_zp_ + c * 4; zp_cache[0] = zp_vec[0]; @@ -210,9 +209,8 @@ class Sm70U4ZpIteratorB { void cache_current_group_metadata(int group) const { CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; if constexpr (kUseMetadataVectorWords) { cache_metadata_vector_words(c, group, cache_n); } else { @@ -237,8 +235,8 @@ class Sm70U4ZpIteratorB { half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data() + frag_base); - marlin::dequant( - static_cast(qword), deq); + marlin::dequant(static_cast(qword), + deq); frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( @@ -259,8 +257,8 @@ class Sm70U4ZpIteratorB { half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data() + frag_base); - marlin::dequant( - static_cast(qword), deq); + marlin::dequant(static_cast(qword), + deq); frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( @@ -271,14 +269,15 @@ class Sm70U4ZpIteratorB { } else { static_assert(ThreadMap::Iterations::kContiguous == 1, "Unsupported SM70 Marlin U4 contiguous iteration count."); - uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_offset_); + uint32_t const qword = + load_qword_vector<1>(qweight_ + qweight_base_offset_); half2 const* scale_vec = cached_scales_; half2 const* zp_vec = cached_zp_; half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data()); - marlin::dequant( - static_cast(qword), deq); + marlin::dequant(static_cast(qword), + deq); frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( @@ -290,17 +289,15 @@ class Sm70U4ZpIteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; if constexpr (kGroupSize != -1) { cache_current_group_metadata(scale_group(logical_k_s)); } int const qweight_base_s = qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - auto const qwords = - load_qword_vector<4>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<4>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -314,20 +311,15 @@ class Sm70U4ZpIteratorB { half2* frag_vec = reinterpret_cast(frag.data() + frag_base); marlin::dequant( static_cast(qword), deq); - frag_vec[0] = - __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); - frag_vec[1] = - __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); + frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); + frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( static_cast(qword >> 8), deq); - frag_vec[2] = - __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); - frag_vec[3] = - __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); + frag_vec[2] = __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); + frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - auto const qwords = - load_qword_vector<2>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<2>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -341,21 +333,19 @@ class Sm70U4ZpIteratorB { half2* frag_vec = reinterpret_cast(frag.data() + frag_base); marlin::dequant( static_cast(qword), deq); - frag_vec[0] = - __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); - frag_vec[1] = - __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); + frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); + frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( static_cast(qword >> 8), deq); - frag_vec[2] = - __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); - frag_vec[3] = - __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); + frag_vec[2] = __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); + frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin U4 contiguous iteration count."); - uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_s); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin U4 contiguous iteration count."); + uint32_t const qword = + load_qword_vector<1>(qweight_ + qweight_base_s); constexpr int kAccess = ThreadMap::kElementsPerAccess; int const frag_base = s * kAccess; half2 const* scale_vec = cached_scales_; @@ -363,8 +353,8 @@ class Sm70U4ZpIteratorB { half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data() + frag_base); - marlin::dequant( - static_cast(qword), deq); + marlin::dequant(static_cast(qword), + deq); frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( @@ -382,8 +372,7 @@ class Sm70U4ZpIteratorB { return; } - if constexpr (kGroupSize != -1 && - ThreadMap::Iterations::kStrided == 1) { + if constexpr (kGroupSize != -1 && ThreadMap::Iterations::kStrided == 1) { int const first_logical_k = k_offset_ + thread_offset_.strided(); cache_current_group_metadata(scale_group(first_logical_k)); } @@ -395,32 +384,30 @@ class Sm70U4ZpIteratorB { template struct Sm70U4ZpGemmSpec { template - using IteratorB = - Sm70U4ZpIteratorB; + using IteratorB = Sm70U4ZpIteratorB; }; -template using Sm70U4ZpGemmTraits = - Sm70MarlinGemmTraits, CtaM, - CtaN, CtaK, Warps, WarpM, WarpN, WarpK, - GroupSize, PackedMacroN>; + Sm70MarlinGemmTraits, CtaM, CtaN, + CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>; -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_u4_gemm_kernel( +__global__ __launch_bounds__(Warps * 32, 1) void sm70_marlin_u4_gemm_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, cutlass::half_t const* __restrict__ b_scales, cutlass::half_t const* __restrict__ b_zeros, cutlass::half_t* __restrict__ c, int m, int n, int k, int lda) { - using Traits = Sm70U4ZpGemmTraits; + using Traits = + Sm70U4ZpGemmTraits; using Mma = typename Traits::Mma; using Epilogue = typename Traits::Epilogue; @@ -440,10 +427,10 @@ void sm70_marlin_u4_gemm_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -469,19 +456,20 @@ void sm70_marlin_u4_gemm_kernel( epilogue(output_op, iterator_D, accumulators, iterator_C); } -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_u4_gemm_splitk_kernel( +__global__ +__launch_bounds__(Warps * 32, 1) void sm70_marlin_u4_gemm_splitk_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, cutlass::half_t const* __restrict__ b_scales, cutlass::half_t const* __restrict__ b_zeros, - cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, int k, int lda, int requested_split_k) { - using Traits = Sm70U4ZpGemmTraits; + cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, + int k, int lda, int requested_split_k) { + using Traits = + Sm70U4ZpGemmTraits; using Mma = typename Traits::Mma; using AtomicEpilogue = Sm70AtomicFp32Epilogue; @@ -492,14 +480,13 @@ void sm70_marlin_u4_gemm_splitk_kernel( int const thread_idx = threadIdx.x; int const warp_idx = cutlass::canonical_warp_idx_sync(); int const lane_idx = threadIdx.x % 32; - Sm70SplitKPartition const partition = - sm70_splitk_partition(k, requested_split_k, int(blockIdx.z)); + Sm70SplitKPartition const partition = sm70_splitk_partition( + k, requested_split_k, int(blockIdx.z)); if (partition.partition_k == 0) { return; } - cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, - partition.k_begin}; + cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, partition.k_begin}; cutlass::MatrixCoord tb_offset_B{partition.k_begin, int(blockIdx.y) * CtaN}; cutlass::MatrixCoord tb_offset_C{int(blockIdx.x) * CtaM, int(blockIdx.y) * CtaN}; @@ -507,10 +494,10 @@ void sm70_marlin_u4_gemm_splitk_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -525,9 +512,8 @@ void sm70_marlin_u4_gemm_splitk_kernel( mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); typename AtomicEpilogue::OutputTileIterator iterator_D( - typename AtomicEpilogue::OutputTileIterator::Params(layout_c), - c, cutlass::MatrixCoord(m, n), - thread_idx, tb_offset_C); + typename AtomicEpilogue::OutputTileIterator::Params(layout_c), c, + cutlass::MatrixCoord(m, n), thread_idx, tb_offset_C); AtomicEpilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx); @@ -536,19 +522,23 @@ void sm70_marlin_u4_gemm_splitk_kernel( } // namespace -template -torch::Tensor launch_sm70_marlin_u4_gemm( - torch::Tensor& a, torch::Tensor& c, torch::Tensor& b_q_weight, - torch::Tensor& b_scales, torch::Tensor& b_zeros, int64_t size_m, - int64_t size_n, int64_t size_k, int requested_split_k) { - auto kernel = sm70_marlin_u4_gemm_kernel; - using SharedStorage = typename Sm70U4ZpGemmTraits< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN, UseMetadataVectorWords>::SharedStorage; +torch::Tensor launch_sm70_marlin_u4_gemm(torch::Tensor& a, torch::Tensor& c, + torch::Tensor& b_q_weight, + torch::Tensor& b_scales, + torch::Tensor& b_zeros, int64_t size_m, + int64_t size_n, int64_t size_k, + int requested_split_k) { + auto kernel = + sm70_marlin_u4_gemm_kernel; + using SharedStorage = + typename Sm70U4ZpGemmTraits::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device()).stream(); @@ -558,8 +548,7 @@ torch::Tensor launch_sm70_marlin_u4_gemm( kernel<<>>( reinterpret_cast(a.data_ptr()), reinterpret_cast(b_q_weight.data_ptr()), - reinterpret_cast( - b_scales.data_ptr()), + reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(b_zeros.data_ptr()), reinterpret_cast(c.data_ptr()), static_cast(size_m), static_cast(size_n), @@ -569,8 +558,8 @@ torch::Tensor launch_sm70_marlin_u4_gemm( } TORCH_CHECK(size_k % int64_t(CtaK) == 0, - "SM70 Marlin uint4 requires K divisible by CTA_K=", - CtaK, " for requested_split_k > 1. Got K=", size_k, + "SM70 Marlin uint4 requires K divisible by CTA_K=", CtaK, + " for requested_split_k > 1. Got K=", size_k, ", requested_split_k=", requested_split_k, "."); auto split_kernel = @@ -583,8 +572,7 @@ torch::Tensor launch_sm70_marlin_u4_gemm( // 0009: fp32 split-K scratch (zeroed for the fp32 atomicAdd), M*N floats. // Only allocated on the split_k>1 path, which the auto-selector uses ONLY at // low batch (M<128), so the scratch is bounded and small. - auto c32 = torch::zeros({size_m, size_n}, - a.options().dtype(at::kFloat)); + auto c32 = torch::zeros({size_m, size_n}, a.options().dtype(at::kFloat)); dim3 grid = sm70_marlin_cta_grid(size_m, size_n, CtaM, CtaN); int const active_split_k = @@ -596,9 +584,9 @@ torch::Tensor launch_sm70_marlin_u4_gemm( reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(b_zeros.data_ptr()), reinterpret_cast(c.data_ptr()), - c32.data_ptr(), - static_cast(size_m), static_cast(size_n), - static_cast(size_k), static_cast(a.stride(0)), requested_split_k); + c32.data_ptr(), static_cast(size_m), static_cast(size_n), + static_cast(size_k), static_cast(a.stride(0)), + requested_split_k); C10_CUDA_KERNEL_LAUNCH_CHECK(); // 0009: single deterministic-width fp32->fp16 narrow of the scratch. @@ -606,8 +594,8 @@ torch::Tensor launch_sm70_marlin_u4_gemm( int const narrow_grid = static_cast((numel + narrow_block - 1) / narrow_block); sm70_marlin_narrow_f32_to_f16<256><<>>( - c32.data_ptr(), - reinterpret_cast(c.data_ptr()), numel); + c32.data_ptr(), reinterpret_cast(c.data_ptr()), + numel); C10_CUDA_KERNEL_LAUNCH_CHECK(); return c; @@ -625,38 +613,44 @@ struct Sm70U4Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - return launch_sm70_marlin_u4_gemm( - a, c, b_q_weight, b_scales, b_zeros, size_m, size_n, size_k, requested_split_k); + a, c, b_q_weight, b_scales, b_zeros, size_m, size_n, size_k, + requested_split_k); } }; -torch::Tensor sm70_marlin_u4_gemm( - torch::Tensor& a, torch::Tensor& c, torch::Tensor& b_q_weight, - torch::Tensor& b_scales, torch::Tensor& b_zeros, int64_t size_m, - int64_t size_n, int64_t size_k, int64_t group_size) { +torch::Tensor sm70_marlin_u4_gemm(torch::Tensor& a, torch::Tensor& c, + torch::Tensor& b_q_weight, + torch::Tensor& b_scales, + torch::Tensor& b_zeros, int64_t size_m, + int64_t size_n, int64_t size_k, + int64_t group_size) { c10::cuda::CUDAGuard device_guard(a.device()); - auto const params = - sm70_marlin_dense_auto_params("uint4", group_size, size_m, size_n, - size_k); + auto const params = sm70_marlin_dense_auto_params("uint4", group_size, size_m, + size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_dense_cta_geometry_supported("SM70 Marlin uint4", geometry); - validate_sm70_marlin_dense_cta_n_alignment("SM70 Marlin uint4", geometry, size_n); + validate_sm70_marlin_dense_cta_geometry_supported("SM70 Marlin uint4", + geometry); + validate_sm70_marlin_dense_cta_n_alignment("SM70 Marlin uint4", geometry, + size_n); if (params.use_metadata_vector_words) { Sm70U4Launcher const launcher{ - a, c, b_q_weight, b_scales, b_zeros, size_m, size_n, - size_k, params.requested_split_k}; + a, c, b_q_weight, + b_scales, b_zeros, size_m, + size_n, size_k, params.requested_split_k}; return dispatch_sm70_marlin_geometry( launcher, geometry, params.packed_macro_n, group_size, "uint4"); } Sm70U4Launcher const launcher{ - a, c, b_q_weight, b_scales, b_zeros, size_m, size_n, - size_k, params.requested_split_k}; + a, c, b_q_weight, + b_scales, b_zeros, size_m, + size_n, size_k, params.requested_split_k}; return dispatch_sm70_marlin_geometry( launcher, geometry, params.packed_macro_n, group_size, "uint4"); } diff --git a/csrc/quantization/marlin/sm70_marlin_u4b8_gemm.cu b/csrc/quantization/marlin/sm70_marlin_u4b8_gemm.cu index cf96b3967a..5de4a3a97e 100644 --- a/csrc/quantization/marlin/sm70_marlin_u4b8_gemm.cu +++ b/csrc/quantization/marlin/sm70_marlin_u4b8_gemm.cu @@ -20,25 +20,25 @@ #include "cutlass/layout/tensor_op_multiplicand_sm70.h" #include "cutlass/transform/threadblock/predicated_tile_iterator.h" -using marlin::sm70::Sm70CtaGeometry; -using marlin::sm70::Sm70AtomicFp16Epilogue; -using marlin::sm70::Sm70AtomicFp32Epilogue; -using marlin::sm70::sm70_marlin_narrow_f32_to_f16; -using marlin::sm70::Sm70MarlinGemmTraits; -using marlin::sm70::Sm70SplitKPartition; -using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; -using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; using marlin::sm70::configure_sm70_dynamic_smem; using marlin::sm70::dispatch_sm70_marlin_geometry; using marlin::sm70::kQuantTileK; using marlin::sm70::kQuantTileN; -using marlin::sm70::sm70_marlin_dense_auto_params; using marlin::sm70::load_qword_vector; using marlin::sm70::qword_from_vector; -using marlin::sm70::sm70_marlin_cta_grid; using marlin::sm70::sm70_active_split_k; +using marlin::sm70::sm70_marlin_cta_grid; +using marlin::sm70::sm70_marlin_dense_auto_params; +using marlin::sm70::sm70_marlin_narrow_f32_to_f16; using marlin::sm70::sm70_splitk_partition; +using marlin::sm70::Sm70AtomicFp16Epilogue; +using marlin::sm70::Sm70AtomicFp32Epilogue; +using marlin::sm70::Sm70CtaGeometry; +using marlin::sm70::Sm70MarlinGemmTraits; +using marlin::sm70::Sm70SplitKPartition; using marlin::sm70::u4_packed_macro_n_qweight_offset_from_logical; +using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; +using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; namespace { @@ -54,12 +54,11 @@ class Sm70U4B8IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + using Fragment = cutlass::Array; static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, "SM70 Marlin U4B8 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, + static_assert(ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, "SM70 Marlin U4B8 IteratorB expects one contiguous iteration " "per 64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, @@ -68,7 +67,8 @@ class Sm70U4B8IteratorB { "SM70 Marlin U4B8 IteratorB expects one packed int4 word per " "access."); static_assert(ThreadMap::Iterations::kStrided >= 1, - "SM70 Marlin U4-family IteratorB expects one or more strided iterations."); + "SM70 Marlin U4-family IteratorB expects one or more strided " + "iterations."); struct Params { int size_n; int aligned_initial_k_step; @@ -148,8 +148,7 @@ class Sm70U4B8IteratorB { if constexpr (kGroupSize == -1) { return 0; } else { - static_assert(kGroupSize == 32 || kGroupSize == 64 || - kGroupSize == 128, + static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, "SM70 Marlin U4B8 supports only group sizes " "-1, 32, 64, and 128."); return logical_k / kGroupSize; @@ -159,8 +158,8 @@ class Sm70U4B8IteratorB { CUTLASS_DEVICE static int qweight_offset_from_logical(Params const& params, int logical_k, int logical_n) { - return u4_packed_macro_n_qweight_offset_from_logical(params.size_n, logical_k, - logical_n); + return u4_packed_macro_n_qweight_offset_from_logical( + params.size_n, logical_k, logical_n); } CUTLASS_DEVICE @@ -190,9 +189,8 @@ class Sm70U4B8IteratorB { void cache_current_group_metadata(int group) const { CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; if constexpr (kUseMetadataVectorWords) { cache_metadata_vector_words(c, group, cache_n); } else { @@ -247,15 +245,17 @@ class Sm70U4B8IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin U4B8 contiguous iteration count."); - uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_offset_); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin U4B8 contiguous iteration count."); + uint32_t const qword = + load_qword_vector<1>(qweight_ + qweight_base_offset_); half2 const* scale_vec = cached_scales_; half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data()); - marlin::dequant( - static_cast(qword), deq); + marlin::dequant(static_cast(qword), + deq); frag_vec[0] = __hmul2(deq[0], scale_vec[0]); frag_vec[1] = __hmul2(deq[1], scale_vec[1]); marlin::dequant( @@ -267,17 +267,15 @@ class Sm70U4B8IteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; if constexpr (kGroupSize != -1) { cache_current_group_metadata(scale_group(logical_k_s)); } int const qweight_base_s = qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - auto const qwords = - load_qword_vector<4>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<4>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -298,8 +296,7 @@ class Sm70U4B8IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - auto const qwords = - load_qword_vector<2>(qweight_ + qweight_base_s); + auto const qwords = load_qword_vector<2>(qweight_ + qweight_base_s); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -320,9 +317,11 @@ class Sm70U4B8IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin U4B8 contiguous iteration count."); - uint32_t const qword = load_qword_vector<1>(qweight_ + qweight_base_s); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin U4B8 contiguous iteration count."); + uint32_t const qword = + load_qword_vector<1>(qweight_ + qweight_base_s); constexpr int kAccess = ThreadMap::kElementsPerAccess; int const frag_base = s * kAccess; half2 const* scale_vec = cached_scales_; @@ -348,8 +347,7 @@ class Sm70U4B8IteratorB { return; } - if constexpr (kGroupSize != -1 && - ThreadMap::Iterations::kStrided == 1) { + if constexpr (kGroupSize != -1 && ThreadMap::Iterations::kStrided == 1) { int const first_logical_k = k_offset_ + thread_offset_.strided(); cache_current_group_metadata(scale_group(first_logical_k)); } @@ -361,31 +359,29 @@ class Sm70U4B8IteratorB { template struct Sm70U4B8GemmSpec { template - using IteratorB = - Sm70U4B8IteratorB; + using IteratorB = Sm70U4B8IteratorB; }; -template using Sm70U4B8GemmTraits = - Sm70MarlinGemmTraits, CtaM, - CtaN, CtaK, Warps, WarpM, WarpN, WarpK, - GroupSize, PackedMacroN>; + Sm70MarlinGemmTraits, CtaM, CtaN, + CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>; -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_u4b8_gemm_kernel( +__global__ __launch_bounds__(Warps * 32, 1) void sm70_marlin_u4b8_gemm_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, cutlass::half_t const* __restrict__ b_scales, cutlass::half_t* __restrict__ c, int m, int n, int k, int lda) { - using Traits = Sm70U4B8GemmTraits; + using Traits = + Sm70U4B8GemmTraits; using Mma = typename Traits::Mma; using Epilogue = typename Traits::Epilogue; @@ -405,10 +401,10 @@ void sm70_marlin_u4b8_gemm_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -433,18 +429,19 @@ void sm70_marlin_u4b8_gemm_kernel( epilogue(output_op, iterator_D, accumulators, iterator_C); } -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_u4b8_gemm_splitk_kernel( +__global__ +__launch_bounds__(Warps * 32, 1) void sm70_marlin_u4b8_gemm_splitk_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, cutlass::half_t const* __restrict__ b_scales, - cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, int k, int lda, int requested_split_k) { - using Traits = Sm70U4B8GemmTraits; + cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, + int k, int lda, int requested_split_k) { + using Traits = + Sm70U4B8GemmTraits; using Mma = typename Traits::Mma; using AtomicEpilogue = Sm70AtomicFp32Epilogue; @@ -455,14 +452,13 @@ void sm70_marlin_u4b8_gemm_splitk_kernel( int const thread_idx = threadIdx.x; int const warp_idx = cutlass::canonical_warp_idx_sync(); int const lane_idx = threadIdx.x % 32; - Sm70SplitKPartition const partition = - sm70_splitk_partition(k, requested_split_k, int(blockIdx.z)); + Sm70SplitKPartition const partition = sm70_splitk_partition( + k, requested_split_k, int(blockIdx.z)); if (partition.partition_k == 0) { return; } - cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, - partition.k_begin}; + cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, partition.k_begin}; cutlass::MatrixCoord tb_offset_B{partition.k_begin, int(blockIdx.y) * CtaN}; cutlass::MatrixCoord tb_offset_C{int(blockIdx.x) * CtaM, int(blockIdx.y) * CtaN}; @@ -470,10 +466,10 @@ void sm70_marlin_u4b8_gemm_splitk_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -487,9 +483,8 @@ void sm70_marlin_u4b8_gemm_splitk_kernel( mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); typename AtomicEpilogue::OutputTileIterator iterator_D( - typename AtomicEpilogue::OutputTileIterator::Params(layout_c), - c, cutlass::MatrixCoord(m, n), - thread_idx, tb_offset_C); + typename AtomicEpilogue::OutputTileIterator::Params(layout_c), c, + cutlass::MatrixCoord(m, n), thread_idx, tb_offset_C); AtomicEpilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx); @@ -498,20 +493,23 @@ void sm70_marlin_u4b8_gemm_splitk_kernel( } // namespace -template -torch::Tensor launch_sm70_marlin_u4b8_gemm( - torch::Tensor& a, torch::Tensor& c, torch::Tensor& b_q_weight, - torch::Tensor& b_scales, int64_t size_m, int64_t size_n, int64_t size_k, - int requested_split_k) { +torch::Tensor launch_sm70_marlin_u4b8_gemm(torch::Tensor& a, torch::Tensor& c, + torch::Tensor& b_q_weight, + torch::Tensor& b_scales, + int64_t size_m, int64_t size_n, + int64_t size_k, + int requested_split_k) { auto kernel = - sm70_marlin_u4b8_gemm_kernel; - using SharedStorage = typename Sm70U4B8GemmTraits< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN, UseMetadataVectorWords>::SharedStorage; + using SharedStorage = + typename Sm70U4B8GemmTraits::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -531,14 +529,13 @@ torch::Tensor launch_sm70_marlin_u4b8_gemm( } TORCH_CHECK(size_k % int64_t(CtaK) == 0, - "SM70 Marlin uint4b8 requires K divisible by CTA_K=", - CtaK, " for requested_split_k > 1. Got K=", size_k, + "SM70 Marlin uint4b8 requires K divisible by CTA_K=", CtaK, + " for requested_split_k > 1. Got K=", size_k, ", requested_split_k=", requested_split_k, "."); auto split_kernel = - sm70_marlin_u4b8_gemm_splitk_kernel; smem_bytes = configure_sm70_dynamic_smem(split_kernel); @@ -546,8 +543,7 @@ torch::Tensor launch_sm70_marlin_u4b8_gemm( // 0009: fp32 split-K scratch (zeroed for the fp32 atomicAdd), M*N floats. // Only allocated on the split_k>1 path, which the auto-selector uses ONLY at // low batch (M<128), so the scratch is bounded and small. - auto c32 = torch::zeros({size_m, size_n}, - a.options().dtype(at::kFloat)); + auto c32 = torch::zeros({size_m, size_n}, a.options().dtype(at::kFloat)); dim3 grid = sm70_marlin_cta_grid(size_m, size_n, CtaM, CtaN); int const active_split_k = @@ -558,9 +554,9 @@ torch::Tensor launch_sm70_marlin_u4b8_gemm( reinterpret_cast(b_q_weight.data_ptr()), reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(c.data_ptr()), - c32.data_ptr(), - static_cast(size_m), static_cast(size_n), - static_cast(size_k), static_cast(a.stride(0)), requested_split_k); + c32.data_ptr(), static_cast(size_m), static_cast(size_n), + static_cast(size_k), static_cast(a.stride(0)), + requested_split_k); C10_CUDA_KERNEL_LAUNCH_CHECK(); // 0009: single deterministic-width fp32->fp16 narrow of the scratch. @@ -568,8 +564,8 @@ torch::Tensor launch_sm70_marlin_u4b8_gemm( int const narrow_grid = static_cast((numel + narrow_block - 1) / narrow_block); sm70_marlin_narrow_f32_to_f16<256><<>>( - c32.data_ptr(), - reinterpret_cast(c.data_ptr()), numel); + c32.data_ptr(), reinterpret_cast(c.data_ptr()), + numel); C10_CUDA_KERNEL_LAUNCH_CHECK(); return c; @@ -586,8 +582,8 @@ struct Sm70U4B8Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { return launch_sm70_marlin_u4b8_gemm const launcher{ - a, c, b_q_weight, b_scales, size_m, size_n, size_k, - params.requested_split_k}; + a, c, b_q_weight, b_scales, + size_m, size_n, size_k, params.requested_split_k}; return dispatch_sm70_marlin_geometry( launcher, geometry, params.packed_macro_n, group_size, "uint4b8"); } Sm70U4B8Launcher const launcher{ - a, c, b_q_weight, b_scales, size_m, size_n, size_k, - params.requested_split_k}; + a, c, b_q_weight, b_scales, + size_m, size_n, size_k, params.requested_split_k}; return dispatch_sm70_marlin_geometry( launcher, geometry, params.packed_macro_n, group_size, "uint4b8"); } diff --git a/csrc/quantization/marlin/sm70_marlin_u8_gemm.cu b/csrc/quantization/marlin/sm70_marlin_u8_gemm.cu index 543325da82..f0cc420ed8 100644 --- a/csrc/quantization/marlin/sm70_marlin_u8_gemm.cu +++ b/csrc/quantization/marlin/sm70_marlin_u8_gemm.cu @@ -20,26 +20,26 @@ #include "cutlass/layout/tensor_op_multiplicand_sm70.h" #include "cutlass/transform/threadblock/predicated_tile_iterator.h" -using marlin::sm70::Sm70CtaGeometry; -using marlin::sm70::Sm70AtomicFp16Epilogue; -using marlin::sm70::Sm70AtomicFp32Epilogue; -using marlin::sm70::sm70_marlin_narrow_f32_to_f16; -using marlin::sm70::Sm70MarlinGemmTraits; -using marlin::sm70::Sm70SplitKPartition; -using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; -using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; using marlin::sm70::configure_sm70_dynamic_smem; using marlin::sm70::dispatch_sm70_marlin_geometry_with_group32; using marlin::sm70::kQuantTileK; using marlin::sm70::kQuantTileN; -using marlin::sm70::sm70_marlin_dense_auto_params; using marlin::sm70::load_qword_vector; using marlin::sm70::qword_from_vector; -using marlin::sm70::sm70_marlin_cta_grid; using marlin::sm70::sm70_active_split_k; +using marlin::sm70::sm70_marlin_cta_grid; +using marlin::sm70::sm70_marlin_dense_auto_params; +using marlin::sm70::sm70_marlin_narrow_f32_to_f16; using marlin::sm70::sm70_splitk_partition; +using marlin::sm70::Sm70AtomicFp16Epilogue; +using marlin::sm70::Sm70AtomicFp32Epilogue; +using marlin::sm70::Sm70CtaGeometry; +using marlin::sm70::Sm70MarlinGemmTraits; +using marlin::sm70::Sm70SplitKPartition; using marlin::sm70::u8_packed_macro_n_qweight_offset_from_logical; using marlin::sm70::u8_packed_macro_n_qweight_word_stride; +using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; +using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; namespace { @@ -55,12 +55,11 @@ class Sm70U8ZpIteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + using Fragment = cutlass::Array; static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, "SM70 Marlin U8 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, + static_assert(ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, "SM70 Marlin U8 IteratorB expects one contiguous iteration per " "64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, @@ -119,7 +118,7 @@ class Sm70U8ZpIteratorB { qweight_base_offset_ = qweight_offset_from_logical(params_, logical_k, logical_n); if constexpr (kGroupSize == -1) { - cache_current_group_metadata(0); + cache_current_group_metadata(0); } } @@ -155,8 +154,7 @@ class Sm70U8ZpIteratorB { if constexpr (kGroupSize == -1) { return 0; } else { - static_assert(kGroupSize == 32 || kGroupSize == 64 || - kGroupSize == 128, + static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, "SM70 Marlin U8 supports only group sizes " "-1, 32, 64, and 128."); return logical_k / kGroupSize; @@ -166,8 +164,8 @@ class Sm70U8ZpIteratorB { CUTLASS_DEVICE static int qweight_offset_from_logical(Params const& params, int logical_k, int logical_n) { - return u8_packed_macro_n_qweight_offset_from_logical(params.size_n, logical_k, - logical_n); + return u8_packed_macro_n_qweight_offset_from_logical( + params.size_n, logical_k, logical_n); } CUTLASS_DEVICE @@ -185,8 +183,8 @@ class Sm70U8ZpIteratorB { scale_cache[2] = scale_vec[2]; scale_cache[3] = scale_vec[3]; - half2 const* zp_vec = reinterpret_cast( - zp_ + group * params_.size_n + cache_n); + half2 const* zp_vec = + reinterpret_cast(zp_ + group * params_.size_n + cache_n); half2* zp_cache = cached_zp_ + c * 4; zp_cache[0] = zp_vec[0]; zp_cache[1] = zp_vec[1]; @@ -205,8 +203,8 @@ class Sm70U8ZpIteratorB { scale_cache[2] = scale_vec[2]; scale_cache[3] = scale_vec[3]; - uint4 const zp_words = *reinterpret_cast( - zp_ + group * params_.size_n + cache_n); + uint4 const zp_words = + *reinterpret_cast(zp_ + group * params_.size_n + cache_n); half2 const* zp_vec = reinterpret_cast(&zp_words); half2* zp_cache = cached_zp_ + c * 4; zp_cache[0] = zp_vec[0]; @@ -219,9 +217,8 @@ class Sm70U8ZpIteratorB { void cache_current_group_metadata(int group) const { CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; if constexpr (kUseMetadataVectorWords) { cache_metadata_vector_words(c, group, cache_n); } else { @@ -295,12 +292,12 @@ class Sm70U8ZpIteratorB { half2 deq[2]; half2* frag_vec = reinterpret_cast(frag.data()); - marlin::dequant( - static_cast(qword0), deq); + marlin::dequant(static_cast(qword0), + deq); frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); - marlin::dequant( - static_cast(qword1), deq); + marlin::dequant(static_cast(qword1), + deq); frag_vec[2] = __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } @@ -308,19 +305,17 @@ class Sm70U8ZpIteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; if constexpr (kGroupSize != -1) { cache_current_group_metadata(scale_group(logical_k_s)); } int const qweight_base_s = qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - uint4 const qwords0 = - load_qword_vector<4>(qweight_ + qweight_base_s); - uint4 const qwords1 = load_qword_vector<4>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint4 const qwords0 = load_qword_vector<4>(qweight_ + qweight_base_s); + uint4 const qwords1 = load_qword_vector<4>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -335,22 +330,17 @@ class Sm70U8ZpIteratorB { half2* frag_vec = reinterpret_cast(frag.data() + frag_base); marlin::dequant( static_cast(qword0), deq); - frag_vec[0] = - __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); - frag_vec[1] = - __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); + frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); + frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( static_cast(qword1), deq); - frag_vec[2] = - __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); - frag_vec[3] = - __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); + frag_vec[2] = __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); + frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - uint2 const qwords0 = - load_qword_vector<2>(qweight_ + qweight_base_s); - uint2 const qwords1 = load_qword_vector<2>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint2 const qwords0 = load_qword_vector<2>(qweight_ + qweight_base_s); + uint2 const qwords1 = load_qword_vector<2>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -365,20 +355,17 @@ class Sm70U8ZpIteratorB { half2* frag_vec = reinterpret_cast(frag.data() + frag_base); marlin::dequant( static_cast(qword0), deq); - frag_vec[0] = - __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); - frag_vec[1] = - __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); + frag_vec[0] = __hfma2(deq[0], scale_vec[0], __hneg2(zp_vec[0])); + frag_vec[1] = __hfma2(deq[1], scale_vec[1], __hneg2(zp_vec[1])); marlin::dequant( static_cast(qword1), deq); - frag_vec[2] = - __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); - frag_vec[3] = - __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); + frag_vec[2] = __hfma2(deq[0], scale_vec[2], __hneg2(zp_vec[2])); + frag_vec[3] = __hfma2(deq[1], scale_vec[3], __hneg2(zp_vec[3])); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin U8 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin U8 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_s); uint32_t const qword1 = load_qword_vector<1>( @@ -409,8 +396,7 @@ class Sm70U8ZpIteratorB { return; } - if constexpr (kGroupSize != -1 && - ThreadMap::Iterations::kStrided == 1) { + if constexpr (kGroupSize != -1 && ThreadMap::Iterations::kStrided == 1) { int const first_logical_k = k_offset_ + thread_offset_.strided(); cache_current_group_metadata(scale_group(first_logical_k)); } @@ -422,32 +408,30 @@ class Sm70U8ZpIteratorB { template struct Sm70U8ZpGemmSpec { template - using IteratorB = - Sm70U8ZpIteratorB; + using IteratorB = Sm70U8ZpIteratorB; }; -template using Sm70U8ZpGemmTraits = - Sm70MarlinGemmTraits, CtaM, - CtaN, CtaK, Warps, WarpM, WarpN, WarpK, - GroupSize, PackedMacroN>; + Sm70MarlinGemmTraits, CtaM, CtaN, + CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>; -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_u8_gemm_kernel( +__global__ __launch_bounds__(Warps * 32, 1) void sm70_marlin_u8_gemm_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, cutlass::half_t const* __restrict__ b_scales, cutlass::half_t const* __restrict__ b_zeros, cutlass::half_t* __restrict__ c, int m, int n, int k, int lda) { - using Traits = Sm70U8ZpGemmTraits; + using Traits = + Sm70U8ZpGemmTraits; using Mma = typename Traits::Mma; using Epilogue = typename Traits::Epilogue; @@ -467,10 +451,10 @@ void sm70_marlin_u8_gemm_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -496,19 +480,20 @@ void sm70_marlin_u8_gemm_kernel( epilogue(output_op, iterator_D, accumulators, iterator_C); } -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_u8_gemm_splitk_kernel( +__global__ +__launch_bounds__(Warps * 32, 1) void sm70_marlin_u8_gemm_splitk_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, cutlass::half_t const* __restrict__ b_scales, cutlass::half_t const* __restrict__ b_zeros, - cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, int k, int lda, int requested_split_k) { - using Traits = Sm70U8ZpGemmTraits; + cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, + int k, int lda, int requested_split_k) { + using Traits = + Sm70U8ZpGemmTraits; using Mma = typename Traits::Mma; using AtomicEpilogue = Sm70AtomicFp32Epilogue; @@ -519,14 +504,13 @@ void sm70_marlin_u8_gemm_splitk_kernel( int const thread_idx = threadIdx.x; int const warp_idx = cutlass::canonical_warp_idx_sync(); int const lane_idx = threadIdx.x % 32; - Sm70SplitKPartition const partition = - sm70_splitk_partition(k, requested_split_k, int(blockIdx.z)); + Sm70SplitKPartition const partition = sm70_splitk_partition( + k, requested_split_k, int(blockIdx.z)); if (partition.partition_k == 0) { return; } - cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, - partition.k_begin}; + cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, partition.k_begin}; cutlass::MatrixCoord tb_offset_B{partition.k_begin, int(blockIdx.y) * CtaN}; cutlass::MatrixCoord tb_offset_C{int(blockIdx.x) * CtaM, int(blockIdx.y) * CtaN}; @@ -534,10 +518,10 @@ void sm70_marlin_u8_gemm_splitk_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -552,9 +536,8 @@ void sm70_marlin_u8_gemm_splitk_kernel( mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); typename AtomicEpilogue::OutputTileIterator iterator_D( - typename AtomicEpilogue::OutputTileIterator::Params(layout_c), - c, cutlass::MatrixCoord(m, n), - thread_idx, tb_offset_C); + typename AtomicEpilogue::OutputTileIterator::Params(layout_c), c, + cutlass::MatrixCoord(m, n), thread_idx, tb_offset_C); AtomicEpilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx); @@ -563,20 +546,23 @@ void sm70_marlin_u8_gemm_splitk_kernel( } // namespace -template -torch::Tensor launch_sm70_marlin_u8_gemm( - torch::Tensor& a, torch::Tensor& c, torch::Tensor& b_q_weight, - torch::Tensor& b_scales, torch::Tensor& b_zeros, int64_t size_m, - int64_t size_n, int64_t size_k, int requested_split_k) { +torch::Tensor launch_sm70_marlin_u8_gemm(torch::Tensor& a, torch::Tensor& c, + torch::Tensor& b_q_weight, + torch::Tensor& b_scales, + torch::Tensor& b_zeros, int64_t size_m, + int64_t size_n, int64_t size_k, + int requested_split_k) { auto kernel = sm70_marlin_u8_gemm_kernel; - using SharedStorage = typename Sm70U8ZpGemmTraits< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN, UseMetadataVectorWords>::SharedStorage; + using SharedStorage = + typename Sm70U8ZpGemmTraits::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -597,8 +583,8 @@ torch::Tensor launch_sm70_marlin_u8_gemm( } TORCH_CHECK(size_k % int64_t(CtaK) == 0, - "SM70 Marlin uint8 requires K divisible by CTA_K=", - CtaK, " for requested_split_k > 1. Got K=", size_k, + "SM70 Marlin uint8 requires K divisible by CTA_K=", CtaK, + " for requested_split_k > 1. Got K=", size_k, ", requested_split_k=", requested_split_k, "."); auto split_kernel = @@ -611,8 +597,7 @@ torch::Tensor launch_sm70_marlin_u8_gemm( // 0009: fp32 split-K scratch (zeroed for the fp32 atomicAdd), M*N floats. // Only allocated on the split_k>1 path, which the auto-selector uses ONLY at // low batch (M<128), so the scratch is bounded and small. - auto c32 = torch::zeros({size_m, size_n}, - a.options().dtype(at::kFloat)); + auto c32 = torch::zeros({size_m, size_n}, a.options().dtype(at::kFloat)); dim3 grid = sm70_marlin_cta_grid(size_m, size_n, CtaM, CtaN); int const active_split_k = @@ -624,9 +609,9 @@ torch::Tensor launch_sm70_marlin_u8_gemm( reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(b_zeros.data_ptr()), reinterpret_cast(c.data_ptr()), - c32.data_ptr(), - static_cast(size_m), static_cast(size_n), - static_cast(size_k), static_cast(a.stride(0)), requested_split_k); + c32.data_ptr(), static_cast(size_m), static_cast(size_n), + static_cast(size_k), static_cast(a.stride(0)), + requested_split_k); C10_CUDA_KERNEL_LAUNCH_CHECK(); // 0009: single deterministic-width fp32->fp16 narrow of the scratch. @@ -634,8 +619,8 @@ torch::Tensor launch_sm70_marlin_u8_gemm( int const narrow_grid = static_cast((numel + narrow_block - 1) / narrow_block); sm70_marlin_narrow_f32_to_f16<256><<>>( - c32.data_ptr(), - reinterpret_cast(c.data_ptr()), numel); + c32.data_ptr(), reinterpret_cast(c.data_ptr()), + numel); C10_CUDA_KERNEL_LAUNCH_CHECK(); return c; @@ -653,13 +638,14 @@ struct Sm70U8Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - return launch_sm70_marlin_u8_gemm( - a, c, b_q_weight, b_scales, b_zeros, size_m, size_n, size_k, requested_split_k); + a, c, b_q_weight, b_scales, b_zeros, size_m, size_n, size_k, + requested_split_k); } }; @@ -671,22 +657,25 @@ torch::Tensor sm70_marlin_u8_gemm(torch::Tensor& a, torch::Tensor& c, int64_t group_size) { c10::cuda::CUDAGuard device_guard(a.device()); - auto const params = - sm70_marlin_dense_auto_params("uint8", group_size, size_m, size_n, - size_k); + auto const params = sm70_marlin_dense_auto_params("uint8", group_size, size_m, + size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_dense_cta_geometry_supported("SM70 Marlin uint8", geometry); - validate_sm70_marlin_dense_cta_n_alignment("SM70 Marlin uint8", geometry, size_n); + validate_sm70_marlin_dense_cta_geometry_supported("SM70 Marlin uint8", + geometry); + validate_sm70_marlin_dense_cta_n_alignment("SM70 Marlin uint8", geometry, + size_n); if (params.use_metadata_vector_words) { Sm70U8Launcher const launcher{ - a, c, b_q_weight, b_scales, b_zeros, size_m, size_n, size_k, - params.requested_split_k}; + a, c, b_q_weight, + b_scales, b_zeros, size_m, + size_n, size_k, params.requested_split_k}; return dispatch_sm70_marlin_geometry_with_group32( launcher, geometry, params.packed_macro_n, group_size, "uint8"); } Sm70U8Launcher const launcher{ - a, c, b_q_weight, b_scales, b_zeros, size_m, size_n, size_k, - params.requested_split_k}; + a, c, b_q_weight, + b_scales, b_zeros, size_m, + size_n, size_k, params.requested_split_k}; return dispatch_sm70_marlin_geometry_with_group32( launcher, geometry, params.packed_macro_n, group_size, "uint8"); } diff --git a/csrc/quantization/marlin/sm70_marlin_u8b128_gemm.cu b/csrc/quantization/marlin/sm70_marlin_u8b128_gemm.cu index e8c312437d..e4b8316c08 100644 --- a/csrc/quantization/marlin/sm70_marlin_u8b128_gemm.cu +++ b/csrc/quantization/marlin/sm70_marlin_u8b128_gemm.cu @@ -20,26 +20,26 @@ #include "cutlass/layout/tensor_op_multiplicand_sm70.h" #include "cutlass/transform/threadblock/predicated_tile_iterator.h" -using marlin::sm70::Sm70CtaGeometry; -using marlin::sm70::Sm70AtomicFp16Epilogue; -using marlin::sm70::Sm70AtomicFp32Epilogue; -using marlin::sm70::sm70_marlin_narrow_f32_to_f16; -using marlin::sm70::Sm70MarlinGemmTraits; -using marlin::sm70::Sm70SplitKPartition; -using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; -using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; using marlin::sm70::configure_sm70_dynamic_smem; using marlin::sm70::dispatch_sm70_marlin_geometry_with_group32; using marlin::sm70::kQuantTileK; using marlin::sm70::kQuantTileN; -using marlin::sm70::sm70_marlin_dense_auto_params; using marlin::sm70::load_qword_vector; using marlin::sm70::qword_from_vector; -using marlin::sm70::sm70_marlin_cta_grid; using marlin::sm70::sm70_active_split_k; +using marlin::sm70::sm70_marlin_cta_grid; +using marlin::sm70::sm70_marlin_dense_auto_params; +using marlin::sm70::sm70_marlin_narrow_f32_to_f16; using marlin::sm70::sm70_splitk_partition; +using marlin::sm70::Sm70AtomicFp16Epilogue; +using marlin::sm70::Sm70AtomicFp32Epilogue; +using marlin::sm70::Sm70CtaGeometry; +using marlin::sm70::Sm70MarlinGemmTraits; +using marlin::sm70::Sm70SplitKPartition; using marlin::sm70::u8_packed_macro_n_qweight_offset_from_logical; using marlin::sm70::u8_packed_macro_n_qweight_word_stride; +using marlin::sm70::validate_sm70_marlin_dense_cta_geometry_supported; +using marlin::sm70::validate_sm70_marlin_dense_cta_n_alignment; namespace { @@ -55,12 +55,12 @@ class Sm70U8B128IteratorB { static int const kPackedMacroN = PackedMacroN_; static constexpr bool kUseMetadataVectorWords = UseMetadataVectorWords_; using Element = cutlass::half_t; - using Fragment = cutlass::Array< - Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; - static_assert(Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, - "SM70 Marlin U8B128 IteratorB expects CTA_N in {64, 128, 256}."); - static_assert(ThreadMap::Iterations::kContiguous == - Shape::kN / kQuantTileN, + using Fragment = cutlass::Array; + static_assert( + Shape::kN == 64 || Shape::kN == 128 || Shape::kN == 256, + "SM70 Marlin U8B128 IteratorB expects CTA_N in {64, 128, 256}."); + static_assert(ThreadMap::Iterations::kContiguous == Shape::kN / kQuantTileN, "SM70 Marlin U8B128 IteratorB expects one contiguous iteration " "per 64-column quant tile."); static_assert(ThreadMap::Delta::kContiguous == kQuantTileN, @@ -152,8 +152,7 @@ class Sm70U8B128IteratorB { if constexpr (kGroupSize == -1) { return 0; } else { - static_assert(kGroupSize == 32 || kGroupSize == 64 || - kGroupSize == 128, + static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, "SM70 Marlin U8B128 supports only group sizes " "-1, 32, 64, and 128."); return logical_k / kGroupSize; @@ -163,13 +162,12 @@ class Sm70U8B128IteratorB { CUTLASS_DEVICE static int qweight_offset_from_logical(Params const& params, int logical_k, int logical_n) { - return u8_packed_macro_n_qweight_offset_from_logical(params.size_n, logical_k, - logical_n); + return u8_packed_macro_n_qweight_offset_from_logical( + params.size_n, logical_k, logical_n); } CUTLASS_DEVICE - static int qweight_word_stride_from_logical(Params const&, - int logical_n) { + static int qweight_word_stride_from_logical(Params const&, int logical_n) { return u8_packed_macro_n_qweight_word_stride(); } @@ -200,9 +198,8 @@ class Sm70U8B128IteratorB { void cache_current_group_metadata(int group) const { CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { - int const cache_n = - n_offset_ + thread_offset_.contiguous() + - c * ThreadMap::Delta::kContiguous; + int const cache_n = n_offset_ + thread_offset_.contiguous() + + c * ThreadMap::Delta::kContiguous; if constexpr (kUseMetadataVectorWords) { cache_metadata_vector_words(c, group, cache_n); } else { @@ -263,8 +260,9 @@ class Sm70U8B128IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin U8B128 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin U8B128 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_offset_); uint32_t const qword1 = load_qword_vector<1>( @@ -286,19 +284,17 @@ class Sm70U8B128IteratorB { int const logical_n_base = n_offset_ + thread_offset_.contiguous(); CUTLASS_PRAGMA_UNROLL for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { - int const logical_k_s = - k_offset_ + thread_offset_.strided() + - s * ThreadMap::Delta::kStrided; + int const logical_k_s = k_offset_ + thread_offset_.strided() + + s * ThreadMap::Delta::kStrided; if constexpr (kGroupSize != -1) { cache_current_group_metadata(scale_group(logical_k_s)); } int const qweight_base_s = qweight_offset_from_logical(params_, logical_k_s, logical_n_base); if constexpr (ThreadMap::Iterations::kContiguous == 4) { - uint4 const qwords0 = - load_qword_vector<4>(qweight_ + qweight_base_s); - uint4 const qwords1 = load_qword_vector<4>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint4 const qwords0 = load_qword_vector<4>(qweight_ + qweight_base_s); + uint4 const qwords1 = load_qword_vector<4>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -320,10 +316,9 @@ class Sm70U8B128IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else if constexpr (ThreadMap::Iterations::kContiguous == 2) { - uint2 const qwords0 = - load_qword_vector<2>(qweight_ + qweight_base_s); - uint2 const qwords1 = load_qword_vector<2>( - qweight_ + qweight_base_s + kQweightWordStrideWords); + uint2 const qwords0 = load_qword_vector<2>(qweight_ + qweight_base_s); + uint2 const qwords1 = load_qword_vector<2>(qweight_ + qweight_base_s + + kQweightWordStrideWords); CUTLASS_PRAGMA_UNROLL for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { constexpr int kAccess = ThreadMap::kElementsPerAccess; @@ -345,8 +340,9 @@ class Sm70U8B128IteratorB { frag_vec[3] = __hmul2(deq[1], scale_vec[3]); } } else { - static_assert(ThreadMap::Iterations::kContiguous == 1, - "Unsupported SM70 Marlin U8B128 contiguous iteration count."); + static_assert( + ThreadMap::Iterations::kContiguous == 1, + "Unsupported SM70 Marlin U8B128 contiguous iteration count."); uint32_t const qword0 = load_qword_vector<1>(qweight_ + qweight_base_s); uint32_t const qword1 = load_qword_vector<1>( @@ -376,8 +372,7 @@ class Sm70U8B128IteratorB { return; } - if constexpr (kGroupSize != -1 && - ThreadMap::Iterations::kStrided == 1) { + if constexpr (kGroupSize != -1 && ThreadMap::Iterations::kStrided == 1) { int const first_logical_k = k_offset_ + thread_offset_.strided(); cache_current_group_metadata(scale_group(first_logical_k)); } @@ -389,32 +384,29 @@ class Sm70U8B128IteratorB { template struct Sm70U8B128GemmSpec { template - using IteratorB = - Sm70U8B128IteratorB; + using IteratorB = Sm70U8B128IteratorB; }; -template using Sm70U8B128GemmTraits = - Sm70MarlinGemmTraits, CtaM, - CtaN, CtaK, Warps, WarpM, WarpN, WarpK, - GroupSize, PackedMacroN>; + Sm70MarlinGemmTraits, CtaM, CtaN, + CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, + PackedMacroN>; -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_u8b128_gemm_kernel( +__global__ __launch_bounds__(Warps * 32, 1) void sm70_marlin_u8b128_gemm_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, cutlass::half_t const* __restrict__ b_scales, cutlass::half_t* __restrict__ c, int m, int n, int k, int lda) { using Traits = Sm70U8B128GemmTraits; + GroupSize, PackedMacroN, UseMetadataVectorWords>; using Mma = typename Traits::Mma; using Epilogue = typename Traits::Epilogue; @@ -434,10 +426,10 @@ void sm70_marlin_u8b128_gemm_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -462,18 +454,19 @@ void sm70_marlin_u8b128_gemm_kernel( epilogue(output_op, iterator_D, accumulators, iterator_C); } -template -__global__ __launch_bounds__(Warps * 32, 1) -void sm70_marlin_u8b128_gemm_splitk_kernel( +__global__ +__launch_bounds__(Warps * 32, 1) void sm70_marlin_u8b128_gemm_splitk_kernel( cutlass::half_t const* __restrict__ a, uint32_t const* __restrict__ b_q_weight, cutlass::half_t const* __restrict__ b_scales, - cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, int k, int lda, int requested_split_k) { - using Traits = Sm70U8B128GemmTraits; + cutlass::half_t* __restrict__ c, float* __restrict__ c32, int m, int n, + int k, int lda, int requested_split_k) { + using Traits = + Sm70U8B128GemmTraits; using Mma = typename Traits::Mma; using AtomicEpilogue = Sm70AtomicFp32Epilogue; @@ -484,14 +477,13 @@ void sm70_marlin_u8b128_gemm_splitk_kernel( int const thread_idx = threadIdx.x; int const warp_idx = cutlass::canonical_warp_idx_sync(); int const lane_idx = threadIdx.x % 32; - Sm70SplitKPartition const partition = - sm70_splitk_partition(k, requested_split_k, int(blockIdx.z)); + Sm70SplitKPartition const partition = sm70_splitk_partition( + k, requested_split_k, int(blockIdx.z)); if (partition.partition_k == 0) { return; } - cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, - partition.k_begin}; + cutlass::MatrixCoord tb_offset_A{int(blockIdx.x) * CtaM, partition.k_begin}; cutlass::MatrixCoord tb_offset_B{partition.k_begin, int(blockIdx.y) * CtaN}; cutlass::MatrixCoord tb_offset_C{int(blockIdx.x) * CtaM, int(blockIdx.y) * CtaN}; @@ -499,10 +491,10 @@ void sm70_marlin_u8b128_gemm_splitk_kernel( typename Traits::LayoutA layout_a(lda); typename Traits::LayoutC layout_c(n); - typename Mma::IteratorA iterator_A( - typename Mma::IteratorA::Params(layout_a), - const_cast(a), cutlass::MatrixCoord(m, k), thread_idx, - tb_offset_A); + typename Mma::IteratorA iterator_A(typename Mma::IteratorA::Params(layout_a), + const_cast(a), + cutlass::MatrixCoord(m, k), thread_idx, + tb_offset_A); typename Mma::IteratorB iterator_B( typename Mma::IteratorB::Params(k, n), reinterpret_cast(b_q_weight), @@ -516,9 +508,8 @@ void sm70_marlin_u8b128_gemm_splitk_kernel( mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); typename AtomicEpilogue::OutputTileIterator iterator_D( - typename AtomicEpilogue::OutputTileIterator::Params(layout_c), - c, cutlass::MatrixCoord(m, n), - thread_idx, tb_offset_C); + typename AtomicEpilogue::OutputTileIterator::Params(layout_c), c, + cutlass::MatrixCoord(m, n), thread_idx, tb_offset_C); AtomicEpilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx); @@ -527,20 +518,23 @@ void sm70_marlin_u8b128_gemm_splitk_kernel( } // namespace -template -torch::Tensor launch_sm70_marlin_u8b128_gemm( - torch::Tensor& a, torch::Tensor& c, torch::Tensor& b_q_weight, - torch::Tensor& b_scales, int64_t size_m, int64_t size_n, int64_t size_k, - int requested_split_k) { +torch::Tensor launch_sm70_marlin_u8b128_gemm(torch::Tensor& a, torch::Tensor& c, + torch::Tensor& b_q_weight, + torch::Tensor& b_scales, + int64_t size_m, int64_t size_n, + int64_t size_k, + int requested_split_k) { auto kernel = sm70_marlin_u8b128_gemm_kernel; - using SharedStorage = typename Sm70U8B128GemmTraits< - CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, - PackedMacroN, UseMetadataVectorWords>::SharedStorage; + using SharedStorage = + typename Sm70U8B128GemmTraits::SharedStorage; size_t smem_bytes = configure_sm70_dynamic_smem(kernel); dim3 block(Warps * 32); @@ -560,23 +554,20 @@ torch::Tensor launch_sm70_marlin_u8b128_gemm( } TORCH_CHECK(size_k % int64_t(CtaK) == 0, - "SM70 Marlin uint8b128 requires K divisible by CTA_K=", - CtaK, " for requested_split_k > 1. Got K=", size_k, + "SM70 Marlin uint8b128 requires K divisible by CTA_K=", CtaK, + " for requested_split_k > 1. Got K=", size_k, ", requested_split_k=", requested_split_k, "."); - auto split_kernel = - sm70_marlin_u8b128_gemm_splitk_kernel; + auto split_kernel = sm70_marlin_u8b128_gemm_splitk_kernel< + CtaM, CtaN, CtaK, Warps, WarpM, WarpN, WarpK, GroupSize, PackedMacroN, + UseMetadataVectorWords>; smem_bytes = configure_sm70_dynamic_smem(split_kernel); int64_t const numel = size_m * size_n; // 0009: fp32 split-K scratch (zeroed for the fp32 atomicAdd), M*N floats. // Only allocated on the split_k>1 path, which the auto-selector uses ONLY at // low batch (M<128), so the scratch is bounded and small. - auto c32 = torch::zeros({size_m, size_n}, - a.options().dtype(at::kFloat)); + auto c32 = torch::zeros({size_m, size_n}, a.options().dtype(at::kFloat)); dim3 grid = sm70_marlin_cta_grid(size_m, size_n, CtaM, CtaN); int const active_split_k = @@ -587,9 +578,9 @@ torch::Tensor launch_sm70_marlin_u8b128_gemm( reinterpret_cast(b_q_weight.data_ptr()), reinterpret_cast(b_scales.data_ptr()), reinterpret_cast(c.data_ptr()), - c32.data_ptr(), - static_cast(size_m), static_cast(size_n), - static_cast(size_k), static_cast(a.stride(0)), requested_split_k); + c32.data_ptr(), static_cast(size_m), static_cast(size_n), + static_cast(size_k), static_cast(a.stride(0)), + requested_split_k); C10_CUDA_KERNEL_LAUNCH_CHECK(); // 0009: single deterministic-width fp32->fp16 narrow of the scratch. @@ -597,8 +588,8 @@ torch::Tensor launch_sm70_marlin_u8b128_gemm( int const narrow_grid = static_cast((numel + narrow_block - 1) / narrow_block); sm70_marlin_narrow_f32_to_f16<256><<>>( - c32.data_ptr(), - reinterpret_cast(c.data_ptr()), numel); + c32.data_ptr(), reinterpret_cast(c.data_ptr()), + numel); C10_CUDA_KERNEL_LAUNCH_CHECK(); return c; @@ -615,12 +606,11 @@ struct Sm70U8B128Launcher { int64_t size_k; int requested_split_k; - template + template torch::Tensor operator()() const { - return launch_sm70_marlin_u8b128_gemm( a, c, b_q_weight, b_scales, size_m, size_n, size_k, requested_split_k); } @@ -633,21 +623,23 @@ torch::Tensor sm70_marlin_u8b128_gemm(torch::Tensor& a, torch::Tensor& c, int64_t group_size) { c10::cuda::CUDAGuard device_guard(a.device()); - auto const params = sm70_marlin_dense_auto_params( - "uint8b128", group_size, size_m, size_n, size_k); + auto const params = sm70_marlin_dense_auto_params("uint8b128", group_size, + size_m, size_n, size_k); Sm70CtaGeometry const geometry = params.geometry; - validate_sm70_marlin_dense_cta_geometry_supported("SM70 Marlin uint8b128", geometry); - validate_sm70_marlin_dense_cta_n_alignment("SM70 Marlin uint8b128", geometry, size_n); + validate_sm70_marlin_dense_cta_geometry_supported("SM70 Marlin uint8b128", + geometry); + validate_sm70_marlin_dense_cta_n_alignment("SM70 Marlin uint8b128", geometry, + size_n); if (params.use_metadata_vector_words) { Sm70U8B128Launcher const launcher{ - a, c, b_q_weight, b_scales, size_m, size_n, size_k, - params.requested_split_k}; + a, c, b_q_weight, b_scales, + size_m, size_n, size_k, params.requested_split_k}; return dispatch_sm70_marlin_geometry_with_group32( launcher, geometry, params.packed_macro_n, group_size, "uint8b128"); } Sm70U8B128Launcher const launcher{ - a, c, b_q_weight, b_scales, size_m, size_n, size_k, - params.requested_split_k}; + a, c, b_q_weight, b_scales, + size_m, size_n, size_k, params.requested_split_k}; return dispatch_sm70_marlin_geometry_with_group32( launcher, geometry, params.packed_macro_n, group_size, "uint8b128"); } diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/gemm.cu b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/gemm.cu index 8fee0cd95f..4c776b38ab 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/gemm.cu +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/gemm.cu @@ -280,6 +280,23 @@ std::optional GetSm70Fp8BlockPrefillPrescaledTarget( return Sm70AwqTp2FastTarget{ desc.n, desc.k, 64, 256, 16, 1, 3, true, "sm70_fp8_pscale_full"}; } + if (desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_1x1536x4096_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 128, 64, 5, 1, true, + "sm70_fp8_pscale_m1"}; + } + if (desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_1x8192x1024_1" || + desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_1x4096x2048_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 128, 64, 2, 3, true, + "sm70_fp8_pscale_m1"}; + } + if (desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_1x1024x4096_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 128, 64, 7, 1, true, + "sm70_fp8_pscale_m1"}; + } + if (desc_str == "sm70_f16_e4m3k128_f16_tnt_fff_1x4096x512_1") { + return Sm70AwqTp2FastTarget{desc.n, desc.k, 8, 128, 64, 2, 1, true, + "sm70_fp8_pscale_m1"}; + } return std::nullopt; } @@ -461,8 +478,13 @@ struct Gemm::Impl { LaunchSpec Dispatch(Context& ctx, DispatchPolicy policy, size_t barriers_size, size_t partials_size) { const auto& desc = ctx.desc(); + const bool allow_prescaled = + policy & DispatchPolicy::kSm70Fp8PrefillPrescaled; const auto is_feasible = [&](const LaunchSpec& spec) { return spec.kernel && + (allow_prescaled || + spec.kernel->name().find("_sm70_fp8_pscale") == + std::string::npos) && spec.kernel->is_feasible(ctx.get_desc(*spec.kernel)); }; if (policy & DispatchPolicy::kSm70Fp8PrefillPrescaled) { @@ -471,7 +493,7 @@ struct Gemm::Impl { return *spec; } if (auto fast_target = GetSm70Fp8BlockPrefillPrescaledTarget(desc)) { - auto specs = Find(ctx, barriers_size, partials_size, 0); + auto specs = Find(ctx, barriers_size, partials_size, 0, true); if (auto fast_spec = SelectSm70AwqTp2FastSpec( ctx, specs, *fast_target, barriers_size, partials_size)) { sm70_fp8_prefill_cache_.Insert(desc, *fast_spec); @@ -482,7 +504,7 @@ struct Gemm::Impl { } if (policy & DispatchPolicy::kMxfp4MoeGroupedM8Fast) { if (auto fast_target = GetSm70Mxfp4MoeGroupedM8FastTarget(desc)) { - auto specs = Find(ctx, barriers_size, partials_size, 0); + auto specs = Find(ctx, barriers_size, partials_size, 0, false); if (auto fast_spec = SelectSm70AwqTp2FastSpec( ctx, specs, *fast_target, barriers_size, partials_size)) { return *fast_spec; @@ -506,7 +528,8 @@ struct Gemm::Impl { } const auto fast_target = GetSm70AwqTp2FastTarget(desc); - auto specs = Find(ctx, barriers_size, partials_size, fast_target ? 0 : 1); + auto specs = + Find(ctx, barriers_size, partials_size, fast_target ? 0 : 1, false); if (!specs.empty()) { auto selected = specs.front(); if (fast_target) { @@ -522,8 +545,16 @@ struct Gemm::Impl { } std::vector Find(Context& ctx, size_t barrier_size, - size_t partials_size, int top_k) { + size_t partials_size, int top_k, + bool include_prescaled) { std::vector feasible = ctx.Filter(registry_.kernels()); + if (!include_prescaled) { + feasible.erase( + std::remove_if(feasible.begin(), feasible.end(), [](const Kernel* k) { + return k->name().find("_sm70_fp8_pscale") != std::string::npos; + }), + feasible.end()); + } std::vector> clusters; { @@ -618,7 +649,8 @@ struct Gemm::Impl { // std::cerr << "GEMM: " << desc.m << "x" << desc.n << "x" << desc.k << // "\n"; - const auto tmp = Find(ctx, barriers_size, partials_size, tuning_.top_k); + const auto tmp = + Find(ctx, barriers_size, partials_size, tuning_.top_k, false); std::vector specs; for (const auto& spec : tmp) { @@ -761,7 +793,7 @@ int Gemm::Run(const Operation& operation, float alpha, const void* A, (preserve_default_kernel || preserve_default_split_count)) { auto default_specs = impl_->Find(*dispatch_context, workspace.barriers_size, - workspace.partials_size, 1); + workspace.partials_size, 1, false); if (!default_specs.empty()) { const auto default_spec = default_specs.front(); if (preserve_default_kernel) { diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_8.cu b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_8.cu index b219e99aec..1a1a0219c2 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_8.cu +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_8.cu @@ -27,6 +27,27 @@ class Qwen38PrescaledFullTileKernelImpl final : public KernelImpl { } }; +template +class Fp8PrescaledM1KernelImpl final : public KernelImpl { + public: + Fp8PrescaledM1KernelImpl() { + this->info_.name += "_sm70_fp8_pscale_m1"; + } + + bool is_feasible(const GemmDesc& desc) const noexcept override { + if (desc.m != 1 || desc.num != 1) { + return false; + } + const bool accepted_shape = + (desc.n == 1536 && desc.k == 4096) || + (desc.n == 8192 && desc.k == 1024) || + (desc.n == 4096 && desc.k == 2048) || + (desc.n == 1024 && desc.k == 4096) || + (desc.n == 4096 && desc.k == 512); + return accepted_shape && KernelImpl::is_feasible(desc); + } +}; + } // namespace void Registry::sm70_884_8() { @@ -50,6 +71,9 @@ void Registry::sm70_884_8() { using CP = Config_E4M3_Prescaled; using Q38PrescaledFull = CP::Type<64, 256, 16, 1, 4, 1, D, S, 2, true, 1, 128, 64, 128, 1, true>; Add(std::make_unique>()); + + using Fp8PrescaledM1 = CP::Type<8, 128, 64, 1, 4, 1, D, S, 2, true, 1, 128>; + Add(std::make_unique>()); // clang-format on } } diff --git a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu index dec4625017..75e0f90e4b 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -1074,7 +1074,7 @@ bool mxfp4_tune_small_shapes_enabled() { bool mxfp4_moe_compact_grouped_decode_enabled() { const char* raw = std::getenv("VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE"); - return raw != nullptr && std::atoi(raw) != 0; + return raw == nullptr || std::atoi(raw) != 0; } bool mxfp4_moe_broadcast_input_decode_enabled() { @@ -1817,6 +1817,290 @@ __device__ __forceinline__ float warp_reduce_sum(float val) { return val; } +__device__ __forceinline__ float warp_group4_sum(float value) { + constexpr unsigned kSinkhornMask = 0x0000ffff; + value += __shfl_xor_sync(kSinkhornMask, value, 1); + value += __shfl_xor_sync(kSinkhornMask, value, 2); + return value; +} + +__device__ __forceinline__ float warp_group4_max(float value) { + constexpr unsigned kSinkhornMask = 0x0000ffff; + value = fmaxf(value, __shfl_xor_sync(kSinkhornMask, value, 1)); + value = fmaxf(value, __shfl_xor_sync(kSinkhornMask, value, 2)); + return value; +} + +__device__ __forceinline__ float warp_stride4_sum(float value) { + constexpr unsigned kSinkhornMask = 0x0000ffff; + value += __shfl_xor_sync(kSinkhornMask, value, 4); + value += __shfl_xor_sync(kSinkhornMask, value, 8); + return value; +} + +__device__ __forceinline__ float sm70_sigmoid(float value) { + return 1.f / (1.f + __expf(-value)); +} + +template +__global__ void sm70_glm_mhc_pre_norm_kernel( + const float* gemm_mul, const float* gemm_sqrsum, const float* hc_scale, + const float* hc_base, const half* residual, float* post_mix, + float* comb_mix, half* layer_input, const half* norm_weight, int num_splits, + float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult, + int sinkhorn_repeat, float norm_eps) { + static_assert(kStreams == 4); + static_assert(kThreads % 32 == 0); + __shared__ float pre_shared[kStreams]; + __shared__ float warp_sums[kThreads / 32]; + __shared__ float norm_scale_shared; + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + if (warp == 0) { + float sqrsum = 0.f; + for (int split = 0; split < num_splits; ++split) { + sqrsum += gemm_sqrsum[split]; + } + const float input_rsqrt = + rsqrtf(sqrsum / static_cast(kStreams * kHiddenSize) + rms_eps); + + if (lane < kStreams) { + float pre_logit = 0.f; + float post_logit = 0.f; + for (int split = 0; split < num_splits; ++split) { + const float* row = gemm_mul + split * 24; + pre_logit += row[lane]; + post_logit += row[kStreams + lane]; + } + const float pre = + sm70_sigmoid(pre_logit * input_rsqrt * hc_scale[0] + hc_base[lane]) + + hc_pre_eps; + pre_shared[lane] = pre; + post_mix[lane] = sm70_sigmoid(post_logit * input_rsqrt * hc_scale[1] + + hc_base[kStreams + lane]) * + hc_post_mult; + } + + if (lane < kStreams * kStreams) { + float logit = 0.f; + for (int split = 0; split < num_splits; ++split) { + logit += gemm_mul[split * 24 + 2 * kStreams + lane]; + } + logit = logit * input_rsqrt * hc_scale[2] + hc_base[2 * kStreams + lane]; + const float row_max = warp_group4_max(logit); + float value = __expf(logit - row_max); + value = value / warp_group4_sum(value) + hc_sinkhorn_eps; + value /= warp_stride4_sum(value) + hc_sinkhorn_eps; + for (int iter = 1; iter < sinkhorn_repeat; ++iter) { + value /= warp_group4_sum(value) + hc_sinkhorn_eps; + value /= warp_stride4_sum(value) + hc_sinkhorn_eps; + } + comb_mix[lane] = value; + } + } + __syncthreads(); + + constexpr int kHiddenPairs = kHiddenSize / 2; + constexpr int kPairsPerThread = kHiddenPairs / kThreads; + static_assert(kHiddenPairs % kThreads == 0); + const half2* residual2 = reinterpret_cast(residual); + half2* layer_input2 = reinterpret_cast(layer_input); + half2 staged_values[kPairsPerThread]; + float local_sqrsum = 0.f; +#pragma unroll + for (int item = 0; item < kPairsPerThread; ++item) { + const int hidden2 = tid + item * kThreads; + float2 value = make_float2(0.f, 0.f); +#pragma unroll + for (int stream = 0; stream < kStreams; ++stream) { + const float2 residual_value = + __half22float2(residual2[stream * kHiddenPairs + hidden2]); + value.x += pre_shared[stream] * residual_value.x; + value.y += pre_shared[stream] * residual_value.y; + } + local_sqrsum += value.x * value.x + value.y * value.y; + staged_values[item] = __floats2half2_rn(value.x, value.y); + } + local_sqrsum = warp_reduce_sum(local_sqrsum); + if (lane == 0) { + warp_sums[warp] = local_sqrsum; + } + __syncthreads(); + if (warp == 0) { + float block_sum = lane < (kThreads / 32) ? warp_sums[lane] : 0.f; + block_sum = warp_reduce_sum(block_sum); + if (lane == 0) { + norm_scale_shared = + rsqrtf(block_sum / static_cast(kHiddenSize) + norm_eps); + } + } + __syncthreads(); + + const float norm_scale = norm_scale_shared; + const half2* norm_weight2 = reinterpret_cast(norm_weight); +#pragma unroll + for (int item = 0; item < kPairsPerThread; ++item) { + const int hidden2 = tid + item * kThreads; + const float2 staged = __half22float2(staged_values[item]); + const float2 weight = __half22float2(norm_weight2[hidden2]); + layer_input2[hidden2] = __floats2half2_rn(staged.x * norm_scale * weight.x, + staged.y * norm_scale * weight.y); + } +} + +template +__global__ void sm70_glm_kda_fg_b_kernel(half* f_out, half* g_out, + const half* f_input, + const half* g_input, + const half* f_weight, + const half* g_weight) { + constexpr int kThreads = kWarps * WARP_SIZE; + constexpr int kCols2 = kCols / 2; + static_assert(kThreads <= 1024); + static_assert(kCols % 2 == 0); + + const int warp = threadIdx.x / WARP_SIZE; + const int lane = threadIdx.x % WARP_SIZE; + const int combined_row = blockIdx.x * kWarps + warp; + if (combined_row >= 2 * kRows) { + return; + } + + const bool is_g = combined_row >= kRows; + const int row = is_g ? combined_row - kRows : combined_row; + const half2* input = reinterpret_cast(is_g ? g_input : f_input); + const half2* weight = reinterpret_cast( + (is_g ? g_weight : f_weight) + row * kCols); + + float dot = 0.f; + for (int col2 = lane; col2 < kCols2; col2 += WARP_SIZE) { + const float2 input_value = __half22float2(input[col2]); + const float2 weight_value = __half22float2(weight[col2]); + dot = fmaf(input_value.x, weight_value.x, dot); + dot = fmaf(input_value.y, weight_value.y, dot); + } + dot = warp_reduce_sum(dot); + if (lane == 0) { + (is_g ? g_out : f_out)[row] = __float2half_rn(dot); + } +} + +void sm70_glm_mhc_pre_norm_out( + torch::Tensor gemm_mul, torch::Tensor gemm_sqrsum, torch::Tensor hc_scale, + torch::Tensor hc_base, torch::Tensor residual, torch::Tensor post_mix, + torch::Tensor comb_mix, torch::Tensor layer_input, + torch::Tensor norm_weight, double rms_eps, double hc_pre_eps, + double hc_sinkhorn_eps, double hc_post_mult, int64_t sinkhorn_repeat, + double norm_eps) { + TORCH_CHECK( + gemm_mul.is_cuda() && gemm_sqrsum.is_cuda() && hc_scale.is_cuda() && + hc_base.is_cuda() && residual.is_cuda() && post_mix.is_cuda() && + comb_mix.is_cuda() && layer_input.is_cuda() && norm_weight.is_cuda(), + "sm70_glm_mhc_pre_norm_out: all tensors must be CUDA."); + TORCH_CHECK(gemm_mul.scalar_type() == torch::kFloat32 && + gemm_sqrsum.scalar_type() == torch::kFloat32 && + hc_scale.scalar_type() == torch::kFloat32 && + hc_base.scalar_type() == torch::kFloat32 && + post_mix.scalar_type() == torch::kFloat32 && + comb_mix.scalar_type() == torch::kFloat32, + "sm70_glm_mhc_pre_norm_out: mix tensors must be float32."); + TORCH_CHECK(residual.scalar_type() == torch::kFloat16 && + layer_input.scalar_type() == torch::kFloat16 && + norm_weight.scalar_type() == torch::kFloat16, + "sm70_glm_mhc_pre_norm_out: activations must be float16."); + TORCH_CHECK(gemm_mul.is_contiguous() && gemm_sqrsum.is_contiguous() && + hc_scale.is_contiguous() && hc_base.is_contiguous() && + residual.is_contiguous() && post_mix.is_contiguous() && + comb_mix.is_contiguous() && layer_input.is_contiguous() && + norm_weight.is_contiguous(), + "sm70_glm_mhc_pre_norm_out: tensors must be contiguous."); + TORCH_CHECK(gemm_mul.dim() == 3, + "sm70_glm_mhc_pre_norm_out: gemm_mul must be rank 3."); + const int64_t num_splits = gemm_mul.size(0); + TORCH_CHECK(gemm_mul.size(1) == 1 && gemm_mul.size(2) == 24 && + gemm_sqrsum.dim() == 2 && gemm_sqrsum.size(0) == num_splits && + gemm_sqrsum.size(1) == 1, + "sm70_glm_mhc_pre_norm_out: invalid staging shape."); + TORCH_CHECK(residual.dim() == 3 && residual.size(0) == 1 && + residual.size(1) == 4 && residual.size(2) == 4096 && + post_mix.numel() == 4 && comb_mix.numel() == 16 && + layer_input.dim() == 2 && layer_input.size(0) == 1 && + layer_input.size(1) == 4096 && norm_weight.numel() == 4096 && + hc_scale.numel() == 3 && hc_base.numel() == 24, + "sm70_glm_mhc_pre_norm_out: invalid GLM mHC shape."); + const auto device = residual.device(); + TORCH_CHECK(gemm_mul.device() == device && gemm_sqrsum.device() == device && + hc_scale.device() == device && hc_base.device() == device && + post_mix.device() == device && comb_mix.device() == device && + layer_input.device() == device && + norm_weight.device() == device, + "sm70_glm_mhc_pre_norm_out: all tensors must share one device."); + TORCH_CHECK(num_splits >= 1 && num_splits <= 8 && sinkhorn_repeat >= 1, + "sm70_glm_mhc_pre_norm_out: invalid reduction configuration."); + const at::cuda::OptionalCUDAGuard device_guard(device_of(residual)); + constexpr int kThreads = 128; + sm70_glm_mhc_pre_norm_kernel<4096, 4, kThreads> + <<<1, kThreads, 0, at::cuda::getCurrentCUDAStream()>>>( + gemm_mul.data_ptr(), gemm_sqrsum.data_ptr(), + hc_scale.data_ptr(), hc_base.data_ptr(), + reinterpret_cast(residual.data_ptr()), + post_mix.data_ptr(), comb_mix.data_ptr(), + reinterpret_cast(layer_input.data_ptr()), + reinterpret_cast(norm_weight.data_ptr()), + static_cast(num_splits), static_cast(rms_eps), + static_cast(hc_pre_eps), static_cast(hc_sinkhorn_eps), + static_cast(hc_post_mult), static_cast(sinkhorn_repeat), + static_cast(norm_eps)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void sm70_glm_kda_fg_b_out(torch::Tensor f_out, torch::Tensor g_out, + torch::Tensor f_input, torch::Tensor g_input, + torch::Tensor f_weight, torch::Tensor g_weight) { + TORCH_CHECK(f_out.is_cuda() && g_out.is_cuda() && f_input.is_cuda() && + g_input.is_cuda() && f_weight.is_cuda() && g_weight.is_cuda(), + "sm70_glm_kda_fg_b_out: all tensors must be CUDA."); + TORCH_CHECK(f_out.scalar_type() == torch::kFloat16 && + g_out.scalar_type() == torch::kFloat16 && + f_input.scalar_type() == torch::kFloat16 && + g_input.scalar_type() == torch::kFloat16 && + f_weight.scalar_type() == torch::kFloat16 && + g_weight.scalar_type() == torch::kFloat16, + "sm70_glm_kda_fg_b_out: all tensors must be float16."); + TORCH_CHECK(f_out.is_contiguous() && g_out.is_contiguous() && + f_input.is_contiguous() && g_input.is_contiguous() && + f_weight.is_contiguous() && g_weight.is_contiguous(), + "sm70_glm_kda_fg_b_out: all tensors must be contiguous."); + TORCH_CHECK(f_out.sizes() == torch::IntArrayRef({1, 2048}) && + g_out.sizes() == torch::IntArrayRef({1, 2048}) && + f_input.sizes() == torch::IntArrayRef({1, 128}) && + g_input.sizes() == torch::IntArrayRef({1, 128}) && + f_weight.sizes() == torch::IntArrayRef({2048, 128}) && + g_weight.sizes() == torch::IntArrayRef({2048, 128}), + "sm70_glm_kda_fg_b_out: requires the GLM TP4 B1 shape."); + const auto device = f_out.device(); + TORCH_CHECK(g_out.device() == device && f_input.device() == device && + g_input.device() == device && f_weight.device() == device && + g_weight.device() == device, + "sm70_glm_kda_fg_b_out: all tensors must share one device."); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(f_out)); + constexpr int kWarps = 8; + constexpr int kThreads = kWarps * WARP_SIZE; + constexpr int kBlocks = (2 * 2048 + kWarps - 1) / kWarps; + sm70_glm_kda_fg_b_kernel<2048, 128, kWarps> + <<>>( + reinterpret_cast(f_out.data_ptr()), + reinterpret_cast(g_out.data_ptr()), + reinterpret_cast(f_input.data_ptr()), + reinterpret_cast(g_input.data_ptr()), + reinterpret_cast(f_weight.data_ptr()), + reinterpret_cast(g_weight.data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + template __global__ void sm70_f16_gate_mul_kernel(half* out, const half* in_feats, const half* gate_weight, int out_rows, @@ -3510,10 +3794,17 @@ void fp8_gemm_sm70_out(torch::Tensor out, torch::Tensor in_feats, TORCH_CHECK(out.stride(1) == 1, "fp8_gemm_sm70: output must be row-major contiguous."); if (exact_8k_prefill_prescaled) { - TORCH_CHECK( - m == 8000 && k == 5120 && (n == 4096 || n == 3584) && !gated_silu, - "fp8_gemm_sm70: pre-scaled block-FP8 prefill requires " - "M=8000, K=5120, N=4096/3584, and no fused epilogue."); + const bool qwen38_prefill = + m == 8000 && k == 5120 && (n == 4096 || n == 3584); + const bool prescaled_m1 = + m == 1 && ((n == 1536 && k == 4096) || (n == 8192 && k == 1024) || + (n == 4096 && k == 2048) || (n == 1024 && k == 4096) || + (n == 4096 && k == 512)); + TORCH_CHECK(qwen38_prefill || prescaled_m1, + "fp8_gemm_sm70: pre-scaled block-FP8 requires an accepted " + "8K prefill or M=1 tensor shape."); + TORCH_CHECK(!gated_silu, + "fp8_gemm_sm70: pre-scaled path does not fuse gated SILU."); } if (gated_silu) { TORCH_CHECK((n % 2) == 0, @@ -5570,6 +5861,17 @@ void fp8_gemm_sm70_prefill_prescaled_out(torch::Tensor out, group_size, k_ld, q_ld, false, true); } +void fp8_gemm_sm70_prescaled_m1_out(torch::Tensor out, torch::Tensor _in_feats, + torch::Tensor _kernel, + torch::Tensor _prescaled_factors, + int64_t group_size, int64_t k_ld, + int64_t q_ld) { + TORCH_CHECK(_in_feats.dim() == 2 && _in_feats.size(0) == 1, + "fp8_gemm_sm70_prescaled_m1_out requires a rank-2 M=1 input."); + vllm::awq_sm70::fp8_gemm_sm70_out(out, _in_feats, _kernel, _prescaled_factors, + group_size, k_ld, q_ld, false, true); +} + void fp8_gemm_sm70_prefill_dispatch_out( torch::Tensor out, int64_t dense_weight_ptr, torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor _scaling_factors, int64_t group_size, @@ -5641,6 +5943,26 @@ void sm70_f16_gemm_out(torch::Tensor out, torch::Tensor _in_feats, vllm::awq_sm70::sm70_f16_gemm_out(out, _in_feats, _kernel, k_ld, gated_silu); } +void sm70_glm_mhc_pre_norm_out( + torch::Tensor gemm_mul, torch::Tensor gemm_sqrsum, torch::Tensor hc_scale, + torch::Tensor hc_base, torch::Tensor residual, torch::Tensor post_mix, + torch::Tensor comb_mix, torch::Tensor layer_input, + torch::Tensor norm_weight, double rms_eps, double hc_pre_eps, + double hc_sinkhorn_eps, double hc_post_mult, int64_t sinkhorn_repeat, + double norm_eps) { + vllm::awq_sm70::sm70_glm_mhc_pre_norm_out( + gemm_mul, gemm_sqrsum, hc_scale, hc_base, residual, post_mix, comb_mix, + layer_input, norm_weight, rms_eps, hc_pre_eps, hc_sinkhorn_eps, + hc_post_mult, sinkhorn_repeat, norm_eps); +} + +void sm70_glm_kda_fg_b_out(torch::Tensor f_out, torch::Tensor g_out, + torch::Tensor f_input, torch::Tensor g_input, + torch::Tensor f_weight, torch::Tensor g_weight) { + vllm::awq_sm70::sm70_glm_kda_fg_b_out(f_out, g_out, f_input, g_input, + f_weight, g_weight); +} + void sm70_f16_indexed_rerank_out(torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _kernel, torch::Tensor candidate_ids, @@ -5808,6 +6130,23 @@ int64_t sm70_gemm_export_cache(torch::Tensor device_hint, #if defined(ENABLE_SM70_TURBOMIND) +namespace { + +__global__ void awq_moe_fill_strided_ptrs_kernel( + turbomind::gemm::StridedPtr* ptrs, char* base, int64_t expert_stride, + int stride, int64_t num_experts) { + const int64_t expert_id = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (expert_id < num_experts) { + ptrs[expert_id] = { + base + expert_id * expert_stride, + stride, + }; + } +} + +} // namespace + std::vector awq_moe_build_strided_ptrs( torch::Tensor tm_weights, // [E, ...] stacked TM weights torch::Tensor tm_scales, // [E, ...] stacked TM scales @@ -5826,12 +6165,6 @@ std::vector awq_moe_build_strided_ptrs( const at::cuda::OptionalCUDAGuard device_guard(device_of(tm_weights)); const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - // Build {ptr, stride} pairs for each expert - std::vector> w_ptrs; - std::vector> s_ptrs; - w_ptrs.reserve(num_experts); - s_ptrs.reserve(num_experts); - const int64_t w_expert_stride = tm_weights.stride(0) * tm_weights.element_size(); const int64_t s_expert_stride = @@ -5839,30 +6172,21 @@ std::vector awq_moe_build_strided_ptrs( char* w_base = static_cast(tm_weights.data_ptr()); char* s_base = static_cast(tm_scales.data_ptr()); - for (int64_t e = 0; e < num_experts; ++e) { - w_ptrs.emplace_back(w_base + e * w_expert_stride, static_cast(k_ld)); - s_ptrs.emplace_back(s_base + e * s_expert_stride, static_cast(q_ld)); - } - - // MakeStridedPtrs allocates GPU memory via cudaMallocAsync - void* w_gpu = turbomind::gemm::MakeStridedPtrs(w_ptrs, stream); - void* s_gpu = turbomind::gemm::MakeStridedPtrs(s_ptrs, stream); - - // Wrap in torch tensors for lifetime management. // StridedPtr is 16 bytes (__align__(16): void* ptr + int stride + padding). const int64_t buf_bytes = num_experts * 16; auto opts = torch::TensorOptions().device(tm_weights.device()).dtype(torch::kUInt8); - - // Copy into torch-managed tensors so cudaFree of the original is safe. auto w_tensor = torch::empty({buf_bytes}, opts); auto s_tensor = torch::empty({buf_bytes}, opts); - cudaMemcpyAsync(w_tensor.data_ptr(), w_gpu, buf_bytes, - cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(s_tensor.data_ptr(), s_gpu, buf_bytes, - cudaMemcpyDeviceToDevice, stream); - cudaFreeAsync(w_gpu, stream); - cudaFreeAsync(s_gpu, stream); + constexpr int kThreads = 256; + const int blocks = static_cast((num_experts + kThreads - 1) / kThreads); + awq_moe_fill_strided_ptrs_kernel<<>>( + reinterpret_cast(w_tensor.data_ptr()), + w_base, w_expert_stride, static_cast(k_ld), num_experts); + awq_moe_fill_strided_ptrs_kernel<<>>( + reinterpret_cast(s_tensor.data_ptr()), + s_base, s_expert_stride, static_cast(q_ld), num_experts); + C10_CUDA_KERNEL_LAUNCH_CHECK(); return {w_tensor, s_tensor}; } @@ -8301,10 +8625,10 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, const bool broadcast_compact_decode_shape = input.size(0) == 1 && out.dim() == 2 && out.size(0) == num_experts && num_experts == 6 && k == 4096 && (n == 512 || n == 1024); - TORCH_CHECK(out.dim() == 2 && out.size(1) == n && - (out.size(0) == input.size(0) || - broadcast_compact_decode_shape), - "mxfp4_moe_dense_stage_sm70_out: out shape mismatch."); + TORCH_CHECK( + out.dim() == 2 && out.size(1) == n && + (out.size(0) == input.size(0) || broadcast_compact_decode_shape), + "mxfp4_moe_dense_stage_sm70_out: out shape mismatch."); TORCH_CHECK(expert_offsets.numel() >= num_experts + 1, "mxfp4_moe_dense_stage_sm70_out: expert_offsets too small."); TORCH_CHECK(dense_expert_ids.numel() >= num_experts, @@ -8323,10 +8647,9 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, if (broadcast_compact_decode_shape || (vllm::awq_sm70::mxfp4_moe_compact_grouped_decode_enabled() && compact_decode_shape)) { - mxfp4_moe_gemm_sm70_out_impl(out, input, expert_offsets, ptrs_w, ptrs_s, - num_experts, k, n, group_size, - dense_expert_ids, true, - broadcast_compact_decode_shape); + mxfp4_moe_gemm_sm70_out_impl( + out, input, expert_offsets, ptrs_w, ptrs_s, num_experts, k, n, + group_size, dense_expert_ids, true, broadcast_compact_decode_shape); return; } const bool grouped_m8_shape = @@ -8566,8 +8889,11 @@ void nvfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, input.size(0) > kNvfp4LegacyCompactGroups && num_experts == 256 && ((k == 2048 && (n == 1024 || n == 512 || n == 256)) || (n == 2048 && (k == 512 || k == 256 || k == 128))); + const bool exact_qwen4_exp_prefill_shape = + input.size(0) > kNvfp4LegacyCompactGroups && num_experts == 512 && + ((k == 2560 && n == 320) || (k == 160 && n == 2560)); if (vllm::awq_sm70::nvfp4_moe_grouped_prefill_enabled() && - exact_qwen36_prefill_shape) { + (exact_qwen36_prefill_shape || exact_qwen4_exp_prefill_shape)) { static std::atomic logged_nvfp4_grouped_prefill{0u}; maybe_log_sm70_moe_route_once( logged_nvfp4_grouped_prefill, diff --git a/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu b/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu index 2dd90101e8..c91adc4449 100644 --- a/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu +++ b/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu @@ -597,8 +597,423 @@ void launch_fp8_qpn8_gated_pair_sm70(const uint8_t* codes, codes, group_scales, input, output, hidden, k, m, channel_scales); } +template +__global__ void fp8_qpn8_split_cta_m1_stage1_sm70_kernel( + const uint8_t* __restrict__ codes, const half* __restrict__ channel_scales, + const half* __restrict__ input, float* __restrict__ partials, int n, + int k) { + const int lane = threadIdx.x; + const int tile = blockIdx.x; + const int split = blockIdx.y; + const int quadpair = (lane >> 2) & 3; + const int row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int groups_k16 = k >> 4; + const int groups_per_split = groups_k16 / SplitK; + const int group_begin = split * groups_per_split; + const uint4* code_ptr = reinterpret_cast(codes) + + static_cast(tile) * groups_k16 * 32 + lane; + const half scale = + __ldg(channel_scales + tile * 32 + qpn8_col_from_lane(lane)); + const half2 scale2 = __halves2half2(scale, scale); + + float accum[NAcc][8]; +#pragma unroll + for (int chain = 0; chain < NAcc; ++chain) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + accum[chain][index] = 0.0f; + } + } + +#pragma unroll 4 + for (int group = group_begin; group < group_begin + groups_per_split; + ++group) { + const uint4 packed = __ldcs(code_ptr + static_cast(group) * 32); + half2 weights[8]; + fp8x8_to_half2x4_fast(make_uint2(packed.x, packed.y), weights); + fp8x8_to_half2x4_fast(make_uint2(packed.z, packed.w), weights + 4); +#pragma unroll + for (int index = 0; index < 8; ++index) { + weights[index] = __hmul2(weights[index], scale2); + } + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (row == 0) { + input01 = *reinterpret_cast(input + group * 16); + input23 = *reinterpret_cast(input + group * 16 + 8); + } + const unsigned* a0 = reinterpret_cast(&input01); + const unsigned* a1 = reinterpret_cast(&input23); + const unsigned* b = reinterpret_cast(weights); + VLLM_SM70_MMA_8N8K4(accum[0], a0[0], a0[1], b[0], b[1]); + VLLM_SM70_MMA_8N8K4(accum[1 % NAcc], a0[2], a0[3], b[2], b[3]); + VLLM_SM70_MMA_8N8K4(accum[2 % NAcc], a1[0], a1[1], b[4], b[5]); + VLLM_SM70_MMA_8N8K4(accum[3 % NAcc], a1[2], a1[3], b[6], b[7]); + } + +#pragma unroll + for (int chain = 1; chain < NAcc; ++chain) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + accum[0][index] += accum[chain][index]; + } + } + + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int index = pair * 4 + offset; + const int output_col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); + const int column = tile * 32 + quadpair * 8 + output_col; + partials[static_cast(split) * n + column] = accum[0][index]; + } + } + } +} + +template +__global__ void fp8_qpn8_split_cta_m1_reduce_sm70_kernel( + half* __restrict__ output, const float* __restrict__ partials, int n) { + const int column = blockIdx.x * blockDim.x + threadIdx.x; + if (column >= n) { + return; + } + float value = 0.0f; +#pragma unroll + for (int split = 0; split < SplitK; ++split) { + value += partials[static_cast(split) * n + column]; + } + output[column] = __float2half(value); +} + +template +__global__ void fp8_qpn8_hc_down_silu_reduce_sm70_kernel( + half* __restrict__ lora, half* __restrict__ injection, + const float* __restrict__ partials, int n) { + const int column = blockIdx.x * blockDim.x + threadIdx.x; + if (column >= 324) { + return; + } + float value = 0.0f; +#pragma unroll + for (int split = 0; split < SplitK; ++split) { + value += partials[static_cast(split) * n + column]; + } + const half rounded = __float2half(value); + if (column < 320) { + const float scaled = __half2float(rounded) * 0.25f; + lora[column] = __float2half(scaled / (1.0f + __expf(-scaled))); + } else { + injection[column - 320] = rounded; + } +} + +template +__global__ void fp8_qpn8_hc_up_gate_mix_sm70_kernel( + const uint8_t* __restrict__ codes, const half* __restrict__ channel_scales, + const half* __restrict__ lora, const half* __restrict__ xn, + half* __restrict__ output) { + constexpr int kHC = 4; + constexpr int kHidden = 2560; + constexpr int kK = 320; + constexpr int kGroupsK16 = kK / 16; + constexpr int kGroupsPerSplit = kGroupsK16 / SplitK; + constexpr int kHiddenTiles = kHidden / 32; + __shared__ float partials[kHC][SplitK][32]; + + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int stream = warp / SplitK; + const int split = warp - stream * SplitK; + const int tile = blockIdx.x + stream * kHiddenTiles; + const int quadpair = (lane >> 2) & 3; + const int row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int group_begin = split * kGroupsPerSplit; + const uint4* code_ptr = reinterpret_cast(codes) + + static_cast(tile) * kGroupsK16 * 32 + lane; + const half scale = + __ldg(channel_scales + tile * 32 + qpn8_col_from_lane(lane)); + const half2 scale2 = __halves2half2(scale, scale); + + float accum[NAcc][8]; +#pragma unroll + for (int chain = 0; chain < NAcc; ++chain) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + accum[chain][index] = 0.0f; + } + } + +#pragma unroll + for (int group = group_begin; group < group_begin + kGroupsPerSplit; + ++group) { + const uint4 packed = __ldcs(code_ptr + static_cast(group) * 32); + half2 weights[8]; + fp8x8_to_half2x4_fast(make_uint2(packed.x, packed.y), weights); + fp8x8_to_half2x4_fast(make_uint2(packed.z, packed.w), weights + 4); +#pragma unroll + for (int index = 0; index < 8; ++index) { + weights[index] = __hmul2(weights[index], scale2); + } + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (row == 0) { + input01 = *reinterpret_cast(lora + group * 16); + input23 = *reinterpret_cast(lora + group * 16 + 8); + } + const unsigned* a0 = reinterpret_cast(&input01); + const unsigned* a1 = reinterpret_cast(&input23); + const unsigned* b = reinterpret_cast(weights); + VLLM_SM70_MMA_8N8K4(accum[0], a0[0], a0[1], b[0], b[1]); + VLLM_SM70_MMA_8N8K4(accum[1 % NAcc], a0[2], a0[3], b[2], b[3]); + VLLM_SM70_MMA_8N8K4(accum[2 % NAcc], a1[0], a1[1], b[4], b[5]); + VLLM_SM70_MMA_8N8K4(accum[3 % NAcc], a1[2], a1[3], b[6], b[7]); + } + +#pragma unroll + for (int chain = 1; chain < NAcc; ++chain) { +#pragma unroll + for (int index = 0; index < 8; ++index) { + accum[0][index] += accum[chain][index]; + } + } + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int index = pair * 4 + offset; + const int output_col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); + partials[stream][split][quadpair * 8 + output_col] = accum[0][index]; + } + } + } + __syncthreads(); + + if (threadIdx.x < 32) { + const int inner = blockIdx.x * 32 + threadIdx.x; + float mixed = 0.0f; +#pragma unroll + for (int hc_stream = 0; hc_stream < kHC; ++hc_stream) { + float gate = 0.0f; +#pragma unroll + for (int k_split = 0; k_split < SplitK; ++k_split) { + gate += partials[hc_stream][k_split][threadIdx.x]; + } + gate = __half2float(__float2half(gate)); + const float normalized = __half2float(xn[hc_stream * kHidden + inner]); + mixed += normalized / (1.0f + __expf(-gate)); + } + output[inner] = __float2half(mixed * 0.25f); + } +} + +__global__ void fp8_qpn8_hc_down_transform_sm70_kernel( + const half* __restrict__ down, half* __restrict__ lora, + half* __restrict__ injection, int m, int down_stride) { + constexpr int kLora = 320; + constexpr int kInjection = 4; + constexpr int kLogicalN = kLora + kInjection; + const int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= m * kLogicalN) { + return; + } + const int row = index / kLogicalN; + const int column = index - row * kLogicalN; + const half value = down[static_cast(row) * down_stride + column]; + if (column < kLora) { + const float scaled = __half2float(value) * 0.25f; + lora[static_cast(row) * kLora + column] = + __float2half(scaled / (1.0f + __expf(-scaled))); + } else { + injection[static_cast(row) * kInjection + column - kLora] = value; + } +} + +__global__ void fp8_qpn8_hc_gate_mix_sm70_kernel(const half* __restrict__ xn, + const half* __restrict__ gate, + half* __restrict__ output, + int m) { + constexpr int kHC = 4; + constexpr int kHidden = 2560; + constexpr int kN = kHC * kHidden; + const int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= m * kHidden) { + return; + } + const int row = index / kHidden; + const int inner = index - row * kHidden; + float mixed = 0.0f; +#pragma unroll + for (int stream = 0; stream < kHC; ++stream) { + const size_t offset = + static_cast(row) * kN + stream * kHidden + inner; + const float gate_value = __half2float(gate[offset]); + const float normalized = __half2float(xn[offset]); + mixed += normalized / (1.0f + __expf(-gate_value)); + } + output[index] = __float2half(mixed * 0.25f); +} + } // namespace +void fp8_qpn8_split_cta_m1_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, + torch::Tensor channel_scales, + torch::Tensor partials, int64_t split_k, + int64_t accumulator_chains) { + TORCH_CHECK(out.is_cuda() && input.is_cuda() && codes.is_cuda() && + channel_scales.is_cuda() && partials.is_cuda(), + "fp8_qpn8_split_cta_m1_sm70_out: tensors must be CUDA"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8 && + channel_scales.scalar_type() == torch::kFloat16 && + partials.scalar_type() == torch::kFloat32, + "fp8_qpn8_split_cta_m1_sm70_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + codes.is_contiguous() && channel_scales.is_contiguous() && + partials.is_contiguous(), + "fp8_qpn8_split_cta_m1_sm70_out: tensors must be contiguous"); + const int64_t n = out.size(1); + const int64_t k = input.size(1); + TORCH_CHECK(out.dim() == 2 && input.dim() == 2 && out.size(0) == 1 && + input.size(0) == 1 && n > 0 && n % 32 == 0 && k > 0 && + k % 16 == 0, + "fp8_qpn8_split_cta_m1_sm70_out: shape mismatch"); + TORCH_CHECK(codes.numel() == n * k && channel_scales.numel() == n && + partials.numel() >= split_k * n, + "fp8_qpn8_split_cta_m1_sm70_out: workspace mismatch"); + TORCH_CHECK(split_k == 8 || split_k == 16 || split_k == 32, + "fp8_qpn8_split_cta_m1_sm70_out: unsupported split_k"); + TORCH_CHECK((k / 16) % split_k == 0, + "fp8_qpn8_split_cta_m1_sm70_out: invalid split_k"); + TORCH_CHECK(accumulator_chains == 1 || accumulator_chains == 2, + "fp8_qpn8_split_cta_m1_sm70_out: invalid accumulator chains"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const dim3 grid(static_cast(n / 32), + static_cast(split_k)); + auto launch = [&](auto split, auto chains) { + constexpr int kSplit = decltype(split)::value; + constexpr int kChains = decltype(chains)::value; + fp8_qpn8_split_cta_m1_stage1_sm70_kernel + <<>>( + codes.data_ptr(), + reinterpret_cast(channel_scales.data_ptr()), + reinterpret_cast(input.data_ptr()), + partials.data_ptr(), static_cast(n), + static_cast(k)); + constexpr int kThreads = 256; + fp8_qpn8_split_cta_m1_reduce_sm70_kernel + <<<(n + kThreads - 1) / kThreads, kThreads, 0, stream>>>( + reinterpret_cast(out.data_ptr()), + partials.data_ptr(), static_cast(n)); + }; + + if (split_k == 8 && accumulator_chains == 1) { + launch(std::integral_constant{}, std::integral_constant{}); + } else if (split_k == 8) { + launch(std::integral_constant{}, std::integral_constant{}); + } else if (split_k == 16 && accumulator_chains == 1) { + launch(std::integral_constant{}, std::integral_constant{}); + } else if (split_k == 16) { + launch(std::integral_constant{}, std::integral_constant{}); + } else if (accumulator_chains == 1) { + launch(std::integral_constant{}, std::integral_constant{}); + } else { + launch(std::integral_constant{}, std::integral_constant{}); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void fp8_qpn8_hc_down_silu_sm70_out(torch::Tensor lora, torch::Tensor injection, + torch::Tensor input, torch::Tensor codes, + torch::Tensor channel_scales, + torch::Tensor partials) { + TORCH_CHECK(lora.is_cuda() && injection.is_cuda() && input.is_cuda() && + codes.is_cuda() && channel_scales.is_cuda() && + partials.is_cuda(), + "fp8_qpn8_hc_down_silu_sm70_out: tensors must be CUDA"); + TORCH_CHECK(lora.scalar_type() == torch::kFloat16 && + injection.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8 && + channel_scales.scalar_type() == torch::kFloat16 && + partials.scalar_type() == torch::kFloat32, + "fp8_qpn8_hc_down_silu_sm70_out: dtype mismatch"); + TORCH_CHECK(lora.is_contiguous() && injection.is_contiguous() && + input.is_contiguous() && codes.is_contiguous() && + channel_scales.is_contiguous() && partials.is_contiguous(), + "fp8_qpn8_hc_down_silu_sm70_out: tensors must be contiguous"); + constexpr int64_t kN = 352; + constexpr int64_t kK = 10240; + constexpr int64_t kSplitK = 32; + TORCH_CHECK(lora.numel() == 320 && injection.numel() == 4 && + input.dim() == 2 && input.size(0) == 1 && + input.size(1) == kK && codes.numel() >= kN * kK && + channel_scales.numel() >= kN && + partials.numel() >= kSplitK * kN, + "fp8_qpn8_hc_down_silu_sm70_out: shape mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + fp8_qpn8_split_cta_m1_stage1_sm70_kernel<32, 1> + <<>>( + codes.data_ptr(), + reinterpret_cast(channel_scales.data_ptr()), + reinterpret_cast(input.data_ptr()), + partials.data_ptr(), kN, kK); + constexpr int kThreads = 256; + fp8_qpn8_hc_down_silu_reduce_sm70_kernel<32><<<2, kThreads, 0, stream>>>( + reinterpret_cast(lora.data_ptr()), + reinterpret_cast(injection.data_ptr()), + partials.data_ptr(), kN); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void fp8_qpn8_hc_up_gate_mix_sm70_out(torch::Tensor out, torch::Tensor lora, + torch::Tensor xn, torch::Tensor codes, + torch::Tensor channel_scales) { + TORCH_CHECK(out.is_cuda() && lora.is_cuda() && xn.is_cuda() && + codes.is_cuda() && channel_scales.is_cuda(), + "fp8_qpn8_hc_up_gate_mix_sm70_out: tensors must be CUDA"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + lora.scalar_type() == torch::kFloat16 && + xn.scalar_type() == torch::kFloat16 && + codes.scalar_type() == torch::kUInt8 && + channel_scales.scalar_type() == torch::kFloat16, + "fp8_qpn8_hc_up_gate_mix_sm70_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && lora.is_contiguous() && + xn.is_contiguous() && codes.is_contiguous() && + channel_scales.is_contiguous(), + "fp8_qpn8_hc_up_gate_mix_sm70_out: tensors must be contiguous"); + constexpr int64_t kHC = 4; + constexpr int64_t kHidden = 2560; + constexpr int64_t kK = 320; + constexpr int64_t kN = kHC * kHidden; + TORCH_CHECK(out.numel() == kHidden && lora.numel() == kK && + xn.numel() == kN && codes.numel() == kN * kK && + channel_scales.numel() == kN, + "fp8_qpn8_hc_up_gate_mix_sm70_out: shape mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(lora)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + fp8_qpn8_hc_up_gate_mix_sm70_kernel<4, 2> + <<>>( + codes.data_ptr(), + reinterpret_cast(channel_scales.data_ptr()), + reinterpret_cast(lora.data_ptr()), + reinterpret_cast(xn.data_ptr()), + reinterpret_cast(out.data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + std::vector fp8_qpn8_prepare_sm70(torch::Tensor qweight, torch::Tensor scales) { TORCH_CHECK(qweight.is_cuda() && scales.is_cuda(), @@ -616,12 +1031,13 @@ std::vector fp8_qpn8_prepare_sm70(torch::Tensor qweight, const int64_t n = qweight.size(0); const int64_t k = qweight.size(1); - TORCH_CHECK(n > 0 && n % 128 == 0 && k > 0 && k % 128 == 0, - "fp8_qpn8_prepare_sm70: N and K must be positive multiples of " - "128"); const bool channel_scales = scales.size(0) == n && scales.size(1) == 1; const bool block_scales = scales.size(0) == n / 128 && scales.size(1) == k / 128; + TORCH_CHECK(channel_scales ? (n > 0 && n % 32 == 0 && k > 0 && k % 16 == 0) + : (n > 0 && n % 128 == 0 && k > 0 && k % 128 == 0), + "fp8_qpn8_prepare_sm70: channel scales require N%32=K%16=0; " + "block scales require N%128=K%128=0"); TORCH_CHECK(channel_scales || block_scales, "fp8_qpn8_prepare_sm70: expected channel scales [N, 1] or " "block scales [N/128, K/128]"); @@ -682,14 +1098,16 @@ void fp8_qpn8_dequantize_sm70_out(torch::Tensor out, torch::Tensor codes, const int64_t k = out.size(0); const int64_t n = out.size(1); - TORCH_CHECK(k > 0 && k % 128 == 0 && n > 0 && n % 128 == 0, - "fp8_qpn8_dequantize_sm70_out: K and N alignment mismatch"); - TORCH_CHECK(codes.numel() == k * n, - "fp8_qpn8_dequantize_sm70_out: packed code size mismatch"); const bool channel_scales = group_scales.size(0) == 1 && group_scales.size(1) == n; const bool block_scales = group_scales.size(0) == k / 128 && group_scales.size(1) == n / 32; + TORCH_CHECK(channel_scales ? (n > 0 && n % 32 == 0 && k > 0 && k % 16 == 0) + : (n > 0 && n % 128 == 0 && k > 0 && k % 128 == 0), + "fp8_qpn8_dequantize_sm70_out: channel scales require " + "N%32=K%16=0; block scales require N%128=K%128=0"); + TORCH_CHECK(codes.numel() == k * n, + "fp8_qpn8_dequantize_sm70_out: packed code size mismatch"); TORCH_CHECK(channel_scales || block_scales, "fp8_qpn8_dequantize_sm70_out: scale shape mismatch"); @@ -776,14 +1194,15 @@ void fp8_qpn8_gemm_sm70_out(torch::Tensor out, torch::Tensor input, TORCH_CHECK(out.size(0) == m, "fp8_qpn8_gemm_sm70_out: output M mismatch"); TORCH_CHECK(n > 0 && n % 32 == 0, "fp8_qpn8_gemm_sm70_out: N must be a positive multiple of 32"); - TORCH_CHECK(k > 0 && k % 128 == 0, - "fp8_qpn8_gemm_sm70_out: K must be a positive multiple of 128"); - TORCH_CHECK(codes.numel() == n * k, - "fp8_qpn8_gemm_sm70_out: packed code size mismatch"); const bool channel_scales = group_scales.size(0) == 1 && group_scales.size(1) == n; const bool block_scales = group_scales.size(0) == k / 128 && group_scales.size(1) == n / 32; + TORCH_CHECK(channel_scales ? (k > 0 && k % 16 == 0) : (k > 0 && k % 128 == 0), + "fp8_qpn8_gemm_sm70_out: channel scales require K%16=0; " + "block scales require K%128=0"); + TORCH_CHECK(codes.numel() == n * k, + "fp8_qpn8_gemm_sm70_out: packed code size mismatch"); TORCH_CHECK(channel_scales || block_scales, "fp8_qpn8_gemm_sm70_out: scale shape mismatch"); TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 12 || split_k == 16 || @@ -926,18 +1345,19 @@ void fp8_qpn8_gemm_ba_split_sm70_out(torch::Tensor qkv_out, torch::Tensor z_out, qkv_out.get_device() == ba_weight.get_device(), "fp8_qpn8_gemm_ba_split_sm70_out: tensors must share one device"); - constexpr int64_t k = 5120; + const int64_t k = input.size(1); constexpr int64_t n = 4096; constexpr int64_t qkv_n = 2560; constexpr int64_t z_n = n - qkv_n; constexpr int64_t ba_n = 24; - TORCH_CHECK(input.dim() == 2 && input.size(0) == 1 && input.size(1) == k, - "fp8_qpn8_gemm_ba_split_sm70_out: expected input [1, 5120]"); + TORCH_CHECK( + input.dim() == 2 && input.size(0) == 1 && (k == 2560 || k == 5120), + "fp8_qpn8_gemm_ba_split_sm70_out: expected K=2560 or 5120"); TORCH_CHECK(qkv_out.numel() == qkv_n && z_out.numel() == z_n && b_out.numel() == ba_n / 2 && a_out.numel() == ba_n / 2, "fp8_qpn8_gemm_ba_split_sm70_out: output shape mismatch"); TORCH_CHECK(codes.dim() == 2 && codes.size(0) == k && codes.size(1) == n, - "fp8_qpn8_gemm_ba_split_sm70_out: expected codes [5120, 4096]"); + "fp8_qpn8_gemm_ba_split_sm70_out: code shape mismatch"); TORCH_CHECK(group_scales.dim() == 2 && group_scales.size(0) == 1 && group_scales.size(1) == n, "fp8_qpn8_gemm_ba_split_sm70_out: expected channel scales"); @@ -972,13 +1392,13 @@ void fp8_qpn8_dispatch_ba_split_sm70_out( torch::Tensor a_out, torch::Tensor qkvz_staging, torch::Tensor ba_staging, int64_t dense_weight_ptr, torch::Tensor input, torch::Tensor codes, torch::Tensor group_scales, torch::Tensor ba_weight) { - constexpr int64_t k = 5120; + const int64_t k = input.size(1); constexpr int64_t n = 4096; constexpr int64_t qkv_n = 2560; constexpr int64_t z_n = n - qkv_n; constexpr int64_t ba_n = 24; const int64_t m = input.size(0); - TORCH_CHECK(input.dim() == 2 && input.size(1) == k && m >= 1, + TORCH_CHECK(input.dim() == 2 && (k == 2560 || k == 5120) && m >= 1, "fp8_qpn8_dispatch_ba_split_sm70_out: bad input shape"); TORCH_CHECK(qkv_out.is_cuda() && z_out.is_cuda() && b_out.is_cuda() && a_out.is_cuda() && qkvz_staging.is_cuda() && @@ -1035,7 +1455,7 @@ void fp8_qpn8_dispatch_ba_split_sm70_out( static std::once_flag route_log_once; std::call_once(route_log_once, []() { std::fprintf(stderr, - "SM70 GDN QPN8 K5120/N4096 + FP16 b/a N24 split C++ route " + "SM70 GDN QPN8 N4096 + FP16 b/a N24 split C++ route " "enabled.\n"); std::fflush(stderr); }); @@ -1091,14 +1511,17 @@ void fp8_qpn8_gated_pair_sm70_out(torch::Tensor out, torch::Tensor input, const int64_t n = hidden * 2; TORCH_CHECK(m >= 1 && m <= 8 && out.size(0) == m, "fp8_qpn8_gated_pair_sm70_out: M must be in [1, 8]"); - TORCH_CHECK(hidden > 0 && hidden % 32 == 0 && k > 0 && k % 128 == 0, - "fp8_qpn8_gated_pair_sm70_out: shape alignment mismatch"); - TORCH_CHECK(codes.numel() == n * k, - "fp8_qpn8_gated_pair_sm70_out: packed code size mismatch"); const bool channel_scales = group_scales.size(0) == 1 && group_scales.size(1) == n; const bool block_scales = group_scales.size(0) == k / 128 && group_scales.size(1) == n / 32; + TORCH_CHECK( + hidden > 0 && hidden % 32 == 0 && + (channel_scales ? (k > 0 && k % 16 == 0) : (k > 0 && k % 128 == 0)), + "fp8_qpn8_gated_pair_sm70_out: channel scales require K%16=0; " + "block scales require K%128=0"); + TORCH_CHECK(codes.numel() == n * k, + "fp8_qpn8_gated_pair_sm70_out: packed code size mismatch"); TORCH_CHECK(channel_scales || block_scales, "fp8_qpn8_gated_pair_sm70_out: scale shape mismatch"); TORCH_CHECK(split_k == 4 || split_k == 8 || split_k == 16, @@ -1203,11 +1626,89 @@ void fp8_qpn8_dispatch_sm70_out(torch::Tensor out, int64_t dense_weight_ptr, gated_silu); } +void fp8_qpn8_hc_dispatch_sm70_out( + torch::Tensor block_out, torch::Tensor injection_out, + torch::Tensor down_staging, torch::Tensor lora_staging, + torch::Tensor gate_staging, torch::Tensor partials, + int64_t dense_weight_ptr, torch::Tensor xn, torch::Tensor down_codes, + torch::Tensor down_scales, torch::Tensor up_codes, + torch::Tensor up_scales) { + constexpr int64_t kHC = 4; + constexpr int64_t kHidden = 2560; + constexpr int64_t kHyperHidden = kHC * kHidden; + constexpr int64_t kLora = 320; + constexpr int64_t kDownComputeN = 352; + constexpr int64_t kDownSplit = 32; + const int64_t m = xn.size(0); + const int64_t kDownN = down_codes.size(1); + TORCH_CHECK(m >= 1 && xn.dim() == 2 && xn.size(1) == kHyperHidden, + "fp8_qpn8_hc_dispatch_sm70_out: expected xn [M, 10240]"); + TORCH_CHECK(block_out.dim() == 2 && block_out.size(0) == m && + block_out.size(1) == kHidden && injection_out.dim() == 2 && + injection_out.size(0) == m && injection_out.size(1) == kHC && + down_staging.dim() == 2 && down_staging.size(0) == m && + down_staging.size(1) == kDownN && lora_staging.dim() == 2 && + lora_staging.size(0) == m && lora_staging.size(1) == kLora && + gate_staging.dim() == 2 && gate_staging.size(0) == m && + gate_staging.size(1) == kHyperHidden, + "fp8_qpn8_hc_dispatch_sm70_out: staging shape mismatch"); + TORCH_CHECK(down_codes.numel() == kHyperHidden * kDownN && + down_scales.numel() == kDownN && + up_codes.numel() == kLora * kHyperHidden && + up_scales.numel() == kHyperHidden && + partials.scalar_type() == torch::kFloat32 && + kDownN >= kDownComputeN && kDownN % 32 == 0 && + partials.numel() >= kDownSplit * kDownComputeN, + "fp8_qpn8_hc_dispatch_sm70_out: weight/workspace mismatch"); + TORCH_CHECK(block_out.scalar_type() == torch::kFloat16 && + injection_out.scalar_type() == torch::kFloat16 && + down_staging.scalar_type() == torch::kFloat16 && + lora_staging.scalar_type() == torch::kFloat16 && + gate_staging.scalar_type() == torch::kFloat16 && + xn.scalar_type() == torch::kFloat16, + "fp8_qpn8_hc_dispatch_sm70_out: activation dtype mismatch"); + + if (m == 1) { + fp8_qpn8_hc_down_silu_sm70_out(lora_staging, injection_out, xn, down_codes, + down_scales, partials); + fp8_qpn8_hc_up_gate_mix_sm70_out(block_out, lora_staging, xn, up_codes, + up_scales); + return; + } + + fp8_qpn8_dispatch_sm70_out(down_staging, dense_weight_ptr, xn, down_codes, + down_scales, 32, 1, false, false); + constexpr int kThreads = 256; + const int down_elements = static_cast(m * (kLora + kHC)); + fp8_qpn8_hc_down_transform_sm70_kernel<<< + (down_elements + kThreads - 1) / kThreads, kThreads, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(down_staging.data_ptr()), + reinterpret_cast(lora_staging.data_ptr()), + reinterpret_cast(injection_out.data_ptr()), + static_cast(m), static_cast(kDownN)); + fp8_qpn8_dispatch_sm70_out(gate_staging, dense_weight_ptr, lora_staging, + up_codes, up_scales, 4, 2, false, false); + const int block_elements = static_cast(m * kHidden); + fp8_qpn8_hc_gate_mix_sm70_kernel<<<(block_elements + kThreads - 1) / kThreads, + kThreads, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(xn.data_ptr()), + reinterpret_cast(gate_staging.data_ptr()), + reinterpret_cast(block_out.data_ptr()), + static_cast(m)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + #ifdef VLLM_QPN8_STANDALONE -// Lets the exact same source file be compiled as an operator-race harness -// before paying for a complete vLLM rebuild. Production builds register the -// operator centrally in torch_bindings.cpp and do not define this macro. + // Lets the exact same source file be compiled as an operator-race harness + // before paying for a complete vLLM rebuild. Production builds register the + // operator centrally in torch_bindings.cpp and do not define this macro. + #ifdef VLLM_QPN8_STANDALONE_QWEN38_NAMESPACE +TORCH_LIBRARY_FRAGMENT(_C_qwen38, ops) { + #else TORCH_LIBRARY_FRAGMENT(_C, ops) { + #endif ops.def("fp8_qpn8_prepare_sm70(Tensor qweight, Tensor scales) -> Tensor[]"); ops.impl("fp8_qpn8_prepare_sm70", torch::kCUDA, &fp8_qpn8_prepare_sm70); ops.def( @@ -1251,5 +1752,30 @@ TORCH_LIBRARY_FRAGMENT(_C, ops) { "int accumulator_chains, bool fast_decoder, bool prefetch_codes) -> ()"); ops.impl("fp8_qpn8_gated_pair_sm70_out", torch::kCUDA, &fp8_qpn8_gated_pair_sm70_out); + ops.def( + "fp8_qpn8_split_cta_m1_sm70_out(Tensor(a!) out, Tensor input, " + "Tensor codes, Tensor channel_scales, Tensor partials, int split_k, " + "int accumulator_chains) -> ()"); + ops.impl("fp8_qpn8_split_cta_m1_sm70_out", torch::kCUDA, + &fp8_qpn8_split_cta_m1_sm70_out); + ops.def( + "fp8_qpn8_hc_down_silu_sm70_out(Tensor(a!) lora, Tensor(b!) " + "injection, Tensor input, Tensor codes, Tensor channel_scales, " + "Tensor partials) -> ()"); + ops.impl("fp8_qpn8_hc_down_silu_sm70_out", torch::kCUDA, + &fp8_qpn8_hc_down_silu_sm70_out); + ops.def( + "fp8_qpn8_hc_up_gate_mix_sm70_out(Tensor(a!) out, Tensor lora, " + "Tensor xn, Tensor codes, Tensor channel_scales) -> ()"); + ops.impl("fp8_qpn8_hc_up_gate_mix_sm70_out", torch::kCUDA, + &fp8_qpn8_hc_up_gate_mix_sm70_out); + ops.def( + "fp8_qpn8_hc_dispatch_sm70_out(Tensor(a!) block_out, Tensor(b!) " + "injection_out, Tensor(c!) down_staging, Tensor(d!) lora_staging, " + "Tensor(e!) gate_staging, Tensor(f!) partials, int dense_weight_ptr, " + "Tensor xn, Tensor down_codes, Tensor down_scales, Tensor up_codes, " + "Tensor up_scales) -> ()"); + ops.impl("fp8_qpn8_hc_dispatch_sm70_out", torch::kCUDA, + &fp8_qpn8_hc_dispatch_sm70_out); } #endif diff --git a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu new file mode 100644 index 0000000000..18574d8add --- /dev/null +++ b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu @@ -0,0 +1,398 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright contributors to the vLLM project + +#include + +#include +#include +#include +#include +#include + +namespace { + +__device__ __forceinline__ void dequant_e2m1x8(unsigned packed, half2 scale, + half2 out[4]) { + constexpr unsigned kSign = 0x80008000u; + constexpr unsigned kExponentMantissa = 0x0e000e00u; + unsigned values[4]; + values[0] = ((packed << 12) & kSign) | ((packed << 9) & kExponentMantissa); + values[1] = ((packed << 8) & kSign) | ((packed << 5) & kExponentMantissa); + values[2] = ((packed << 4) & kSign) | ((packed << 1) & kExponentMantissa); + values[3] = (packed & kSign) | ((packed >> 3) & kExponentMantissa); +#pragma unroll + for (int index = 0; index < 4; ++index) { + out[index] = + __hmul2(*reinterpret_cast(&values[index]), scale); + } +} + +#define VLLM_SM70_MMA_8N8K4(C, A0, A1, B0, B1) \ + asm volatile( \ + "mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 " \ + "{%0,%1,%2,%3,%4,%5,%6,%7}, {%8,%9}, {%10,%11}, " \ + "{%0,%1,%2,%3,%4,%5,%6,%7};\n" \ + : "+f"(C[0]), "+f"(C[1]), "+f"(C[2]), "+f"(C[3]), "+f"(C[4]), \ + "+f"(C[5]), "+f"(C[6]), "+f"(C[7]) \ + : "r"(A0), "r"(A1), "r"(B0), "r"(B1)) + +template +__global__ void mxfp4_qpn_m1_sm70_kernel(const half* __restrict__ input, + const uint32_t* __restrict__ weights, + const uint8_t* __restrict__ scales, + const int32_t* __restrict__ expert_ids, + half* __restrict__ output, int n, + int k, bool broadcast_input) { + __shared__ float partials[kSplitK][32]; + + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int tile = blockIdx.x; + const int route = blockIdx.y; + const int expert = __ldg(expert_ids + route); + if (expert < 0 || expert >= 256) { + if (threadIdx.x < 32) { + output[static_cast(route) * n + tile * 32 + threadIdx.x] = + __float2half(0.0f); + } + return; + } + + const int quadpair = (lane >> 2) & 3; + const int a_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int packed_col = + ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) ? 4 : 0); + const int groups_k16 = k >> 4; + const int groups_per_warp = groups_k16 / kSplitK; + const int group_begin = warp * groups_per_warp; + const int groups_k8 = k >> 3; + const int tiles_n32 = n >> 5; + + const size_t words_per_expert = static_cast(k) * n / 8; + const uint32_t* expert_weights = + weights + static_cast(expert) * words_per_expert; + const size_t scales_per_expert = static_cast(k >> 5) * n; + const uint8_t* expert_scales = + scales + static_cast(expert) * scales_per_expert; + const half* input_row = + input + static_cast(broadcast_input ? 0 : route) * k; + + float accum[8] = {}; + const half2 exponent_rebias = __float2half2_rn(16384.0f); +#pragma unroll 4 + for (int group = group_begin; group < group_begin + groups_per_warp; + ++group) { + const size_t tile_group_base = + (static_cast(tile) * groups_k8 + group * 2) * 32 + packed_col; + const unsigned packed0 = __ldcs(expert_weights + tile_group_base); + const unsigned packed1 = __ldcs(expert_weights + tile_group_base + 32); + const size_t scale_index = + (static_cast(group >> 1) * tiles_n32 + tile) * 32 + packed_col; + const uint8_t adjusted_exponent = __ldg(expert_scales + scale_index); + const half scalar = + __ushort_as_half(static_cast(adjusted_exponent) << 10); + const half2 scale = + __hmul2(__halves2half2(scalar, scalar), exponent_rebias); + + half2 decoded[8]; + dequant_e2m1x8(packed0, scale, decoded); + dequant_e2m1x8(packed1, scale, decoded + 4); + const unsigned* b = reinterpret_cast(decoded); + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (a_row == 0) { + input01 = *reinterpret_cast(input_row + group * 16); + input23 = *reinterpret_cast(input_row + group * 16 + 8); + } + const unsigned* a0 = reinterpret_cast(&input01); + const unsigned* a1 = reinterpret_cast(&input23); + VLLM_SM70_MMA_8N8K4(accum, a0[0], a0[1], b[0], b[1]); + VLLM_SM70_MMA_8N8K4(accum, a0[2], a0[3], b[2], b[3]); + VLLM_SM70_MMA_8N8K4(accum, a1[0], a1[1], b[4], b[5]); + VLLM_SM70_MMA_8N8K4(accum, a1[2], a1[3], b[6], b[7]); + } + + // Only eight lane roles own M=1's 32 output columns. + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int index = pair * 4 + offset; + const int local_col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); + partials[warp][quadpair * 8 + local_col] = accum[index]; + } + } + } + __syncthreads(); + + if (warp == 0) { + float value = 0.0f; +#pragma unroll + for (int k_warp = 0; k_warp < kSplitK; ++k_warp) { + value += partials[k_warp][lane]; + } + output[static_cast(route) * n + tile * 32 + lane] = + __float2half(value); + } +} + +template +__global__ void nvfp4_qpn_m1_sm70_kernel(const half* __restrict__ input, + const uint32_t* __restrict__ weights, + const half* __restrict__ scales, + const int32_t* __restrict__ expert_ids, + half* __restrict__ output, int n, + int k, bool broadcast_input) { + __shared__ float partials[kSplitK][32]; + + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int tile = blockIdx.x; + const int route = blockIdx.y; + const int expert = __ldg(expert_ids + route); + if (expert < 0 || expert >= 512) { + if (threadIdx.x < 32) { + output[static_cast(route) * n + tile * 32 + threadIdx.x] = + __float2half(0.0f); + } + return; + } + + const int quadpair = (lane >> 2) & 3; + const int a_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int packed_col = + ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) ? 4 : 0); + const int groups_k16 = k >> 4; + const int groups_per_warp = groups_k16 / kSplitK; + const int group_begin = warp * groups_per_warp; + const int groups_k8 = k >> 3; + const int tiles_n32 = n >> 5; + + const size_t words_per_expert = static_cast(k) * n / 8; + const uint32_t* expert_weights = + weights + static_cast(expert) * words_per_expert; + const size_t scales_per_expert = static_cast(k >> 4) * n; + const half* expert_scales = + scales + static_cast(expert) * scales_per_expert; + const half* input_row = + input + static_cast(broadcast_input ? 0 : route) * k; + + float accum[8] = {}; +#pragma unroll 4 + for (int group = group_begin; group < group_begin + groups_per_warp; + ++group) { + const size_t tile_group_base = + (static_cast(tile) * groups_k8 + group * 2) * 32 + packed_col; + const unsigned packed0 = __ldcs(expert_weights + tile_group_base); + const unsigned packed1 = __ldcs(expert_weights + tile_group_base + 32); + const size_t scale_index = + (static_cast(group) * tiles_n32 + tile) * 32 + packed_col; + const half scalar = __ldg(expert_scales + scale_index); + // dequant_e2m1x8 materializes the FP4 payload with a 2^-14 exponent + // offset so that every E2M1 value can be formed with integer bit ops. + // MXFP4 folds the matching 2^14 correction into its exponent scale; + // NVFP4's prepared FP16 group scale needs the same correction here. + const half2 scale = + __hmul2(__halves2half2(scalar, scalar), __float2half2_rn(16384.0f)); + + half2 decoded[8]; + dequant_e2m1x8(packed0, scale, decoded); + dequant_e2m1x8(packed1, scale, decoded + 4); + const unsigned* b = reinterpret_cast(decoded); + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (a_row == 0) { + input01 = *reinterpret_cast(input_row + group * 16); + input23 = *reinterpret_cast(input_row + group * 16 + 8); + } + const unsigned* a0 = reinterpret_cast(&input01); + const unsigned* a1 = reinterpret_cast(&input23); + VLLM_SM70_MMA_8N8K4(accum, a0[0], a0[1], b[0], b[1]); + VLLM_SM70_MMA_8N8K4(accum, a0[2], a0[3], b[2], b[3]); + VLLM_SM70_MMA_8N8K4(accum, a1[0], a1[1], b[4], b[5]); + VLLM_SM70_MMA_8N8K4(accum, a1[2], a1[3], b[6], b[7]); + } + + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int index = pair * 4 + offset; + const int local_col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); + partials[warp][quadpair * 8 + local_col] = accum[index]; + } + } + } + __syncthreads(); + + if (warp == 0) { + float value = 0.0f; +#pragma unroll + for (int k_warp = 0; k_warp < kSplitK; ++k_warp) { + value += partials[k_warp][lane]; + } + output[static_cast(route) * n + tile * 32 + lane] = + __float2half(value); + } +} + +template +void launch_mxfp4_qpn_m1(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, torch::Tensor scales, + torch::Tensor expert_ids, bool broadcast_input) { + const int n = static_cast(out.size(1)); + const int k = static_cast(input.size(1)); + mxfp4_qpn_m1_sm70_kernel + <<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(weights.data_ptr()), + scales.data_ptr(), expert_ids.data_ptr(), + reinterpret_cast(out.data_ptr()), n, k, + broadcast_input); +} + +template +void launch_nvfp4_qpn_m1(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, torch::Tensor scales, + torch::Tensor expert_ids, bool broadcast_input) { + const int n = static_cast(out.size(1)); + const int k = static_cast(input.size(1)); + nvfp4_qpn_m1_sm70_kernel + <<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(weights.data_ptr()), + reinterpret_cast(scales.data_ptr()), + expert_ids.data_ptr(), + reinterpret_cast(out.data_ptr()), n, k, + broadcast_input); +} + +void dispatch_nvfp4_qpn_m1(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, torch::Tensor scales, + torch::Tensor expert_ids, bool broadcast_input, + int64_t split_k) { +#define VLLM_NVFP4_QPN_M1_CASE(SPLIT) \ + case SPLIT: \ + launch_nvfp4_qpn_m1(out, input, weights, scales, expert_ids, \ + broadcast_input); \ + break + switch (split_k) { + VLLM_NVFP4_QPN_M1_CASE(1); + VLLM_NVFP4_QPN_M1_CASE(2); + VLLM_NVFP4_QPN_M1_CASE(4); + VLLM_NVFP4_QPN_M1_CASE(5); + VLLM_NVFP4_QPN_M1_CASE(8); + VLLM_NVFP4_QPN_M1_CASE(10); + VLLM_NVFP4_QPN_M1_CASE(16); + VLLM_NVFP4_QPN_M1_CASE(20); + VLLM_NVFP4_QPN_M1_CASE(32); + default: + TORCH_CHECK(false, "nvfp4_moe_qpn_m1_sm70_out: unsupported split_k ", + split_k); + } +#undef VLLM_NVFP4_QPN_M1_CASE +} + +} // namespace + +void mxfp4_moe_qpn_m1_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, torch::Tensor scales, + torch::Tensor expert_ids, bool broadcast_input) { + TORCH_CHECK(out.is_cuda() && input.is_cuda() && weights.is_cuda() && + scales.is_cuda() && expert_ids.is_cuda(), + "mxfp4_moe_qpn_m1_sm70_out: tensors must be CUDA tensors"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + weights.scalar_type() == torch::kInt32 && + scales.scalar_type() == torch::kUInt8 && + expert_ids.scalar_type() == torch::kInt32, + "mxfp4_moe_qpn_m1_sm70_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + weights.is_contiguous() && scales.is_contiguous() && + expert_ids.is_contiguous(), + "mxfp4_moe_qpn_m1_sm70_out: tensors must be contiguous"); + TORCH_CHECK(out.dim() == 2 && input.dim() == 2 && weights.dim() == 3 && + scales.dim() == 3 && expert_ids.dim() == 1, + "mxfp4_moe_qpn_m1_sm70_out: rank mismatch"); + TORCH_CHECK(out.size(0) == 6 && expert_ids.numel() == 6 && + weights.size(0) == 256 && scales.size(0) == 256, + "mxfp4_moe_qpn_m1_sm70_out: expected six routes and 256 experts"); + TORCH_CHECK(input.get_device() == out.get_device() && + input.get_device() == weights.get_device() && + input.get_device() == scales.get_device() && + input.get_device() == expert_ids.get_device(), + "mxfp4_moe_qpn_m1_sm70_out: device mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + if (broadcast_input) { + TORCH_CHECK(input.sizes() == torch::IntArrayRef({1, 4096}) && + out.sizes() == torch::IntArrayRef({6, 1024}) && + weights.sizes() == torch::IntArrayRef({256, 4096, 128}) && + scales.sizes() == torch::IntArrayRef({256, 128, 1024}), + "mxfp4_moe_qpn_m1_sm70_out: W13 tensor contract mismatch"); + launch_mxfp4_qpn_m1<16>(out, input, weights, scales, expert_ids, true); + } else { + TORCH_CHECK(input.sizes() == torch::IntArrayRef({6, 512}) && + out.sizes() == torch::IntArrayRef({6, 4096}) && + weights.sizes() == torch::IntArrayRef({256, 512, 512}) && + scales.sizes() == torch::IntArrayRef({256, 16, 4096}), + "mxfp4_moe_qpn_m1_sm70_out: W2 tensor contract mismatch"); + launch_mxfp4_qpn_m1<8>(out, input, weights, scales, expert_ids, false); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void nvfp4_moe_qpn_m1_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, torch::Tensor scales, + torch::Tensor expert_ids, bool broadcast_input, + int64_t split_k) { + TORCH_CHECK(out.is_cuda() && input.is_cuda() && weights.is_cuda() && + scales.is_cuda() && expert_ids.is_cuda(), + "nvfp4_moe_qpn_m1_sm70_out: tensors must be CUDA tensors"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + weights.scalar_type() == torch::kInt32 && + scales.scalar_type() == torch::kFloat16 && + expert_ids.scalar_type() == torch::kInt32, + "nvfp4_moe_qpn_m1_sm70_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + weights.is_contiguous() && scales.is_contiguous() && + expert_ids.is_contiguous(), + "nvfp4_moe_qpn_m1_sm70_out: tensors must be contiguous"); + TORCH_CHECK(out.dim() == 2 && input.dim() == 2 && weights.dim() == 3 && + scales.dim() == 3 && expert_ids.dim() == 1, + "nvfp4_moe_qpn_m1_sm70_out: rank mismatch"); + TORCH_CHECK(out.size(0) == 10 && expert_ids.numel() == 10 && + weights.size(0) == 512 && scales.size(0) == 512, + "nvfp4_moe_qpn_m1_sm70_out: expected ten routes and 512 experts"); + TORCH_CHECK(input.get_device() == out.get_device() && + input.get_device() == weights.get_device() && + input.get_device() == scales.get_device() && + input.get_device() == expert_ids.get_device(), + "nvfp4_moe_qpn_m1_sm70_out: device mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const int64_t groups_k16 = input.size(1) / 16; + TORCH_CHECK(split_k > 0 && split_k <= 32 && groups_k16 % split_k == 0, + "nvfp4_moe_qpn_m1_sm70_out: split_k must divide K/16"); + if (broadcast_input) { + TORCH_CHECK(input.sizes() == torch::IntArrayRef({1, 2560}) && + out.sizes() == torch::IntArrayRef({10, 320}) && + weights.sizes() == torch::IntArrayRef({512, 2560, 40}) && + scales.sizes() == torch::IntArrayRef({512, 160, 320}), + "nvfp4_moe_qpn_m1_sm70_out: W13 tensor contract mismatch"); + } else { + TORCH_CHECK(input.sizes() == torch::IntArrayRef({10, 160}) && + out.sizes() == torch::IntArrayRef({10, 2560}) && + weights.sizes() == torch::IntArrayRef({512, 160, 320}) && + scales.sizes() == torch::IntArrayRef({512, 10, 2560}), + "nvfp4_moe_qpn_m1_sm70_out: W2 tensor contract mismatch"); + } + dispatch_nvfp4_qpn_m1(out, input, weights, scales, expert_ids, + broadcast_input, split_k); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/csrc/sm70_turbomind/ops/tm_registry_sm70.cu b/csrc/sm70_turbomind/ops/tm_registry_sm70.cu index f4102ef0e3..7d692f848f 100644 --- a/csrc/sm70_turbomind/ops/tm_registry_sm70.cu +++ b/csrc/sm70_turbomind/ops/tm_registry_sm70.cu @@ -1,39 +1,38 @@ -// Copyright (c) OpenMMLab. All rights reserved. - -#include "src/turbomind/kernels/gemm/arch.h" -#include "src/turbomind/kernels/gemm/registry.h" -#include "src/turbomind/kernels/gemm/utils.h" - -namespace turbomind::gemm { - -Registry::Registry(std::shared_ptr device_prop): - device_prop_{std::move(device_prop)}, arch_{device_prop_->major * 100 + device_prop_->minor * 10} -{ - // Register the V100 kernels we actually use in this build: - // AWQ uint4, FP8/E4M3, and dense fp16 Tensor Core paths. - sm70_884_4(); - sm70_884_8(); - sm70_884_16(); +// Copyright (c) OpenMMLab. All rights reserved. + +#include "src/turbomind/kernels/gemm/arch.h" +#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/utils.h" + +namespace turbomind::gemm { + +Registry::Registry(std::shared_ptr device_prop) + : device_prop_{std::move(device_prop)}, + arch_{device_prop_->major * 100 + device_prop_->minor * 10} { + // Register the V100 kernels we actually use in this build: + // AWQ uint4, FP8/E4M3, and dense fp16 Tensor Core paths. + sm70_884_4(); + sm70_884_8(); + sm70_884_16(); } - -bool Registry::Add(std::unique_ptr kernel) -{ - bool is_valid = true; - - if (!is_arch_compatible(kernel->arch(), arch_)) { - is_valid = false; - } - - if ((int)device_prop_->sharedMemPerBlockOptin < kernel->smem_size()) { - is_valid = false; - } - - if (is_valid) { - ptrs_.push_back(kernels_.emplace_back(transpose(*kernel)).get()); - ptrs_.push_back(kernels_.emplace_back(std::move(kernel)).get()); - } - - return true; -} - -} // namespace turbomind::gemm + +bool Registry::Add(std::unique_ptr kernel) { + bool is_valid = true; + + if (!is_arch_compatible(kernel->arch(), arch_)) { + is_valid = false; + } + + if ((int)device_prop_->sharedMemPerBlockOptin < kernel->smem_size()) { + is_valid = false; + } + + if (is_valid) { + ptrs_.push_back(kernels_.emplace_back(transpose(*kernel)).get()); + ptrs_.push_back(kernels_.emplace_back(std::move(kernel)).get()); + } + + return true; +} + +} // namespace turbomind::gemm diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 3cdfd346fc..fab3d7e0ef 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -287,6 +287,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("fp8_qpn8_gated_pair_sm70_out", torch::kCUDA, &fp8_qpn8_gated_pair_sm70_out); + ops.def( + "fp8_qpn8_hc_dispatch_sm70_out(Tensor(a!) block_out, Tensor(b!) " + "injection_out, Tensor(c!) down_staging, Tensor(d!) lora_staging, " + "Tensor(e!) gate_staging, Tensor(f!) partials, int dense_weight_ptr, " + "Tensor xn, Tensor down_codes, Tensor down_scales, Tensor up_codes, " + "Tensor up_scales) -> ()"); + ops.impl("fp8_qpn8_hc_dispatch_sm70_out", torch::kCUDA, + &fp8_qpn8_hc_dispatch_sm70_out); + ops.def("nvfp4_qpn4_prepare_sm70(Tensor qweight, Tensor scales) -> Tensor[]"); ops.impl("nvfp4_qpn4_prepare_sm70", torch::kCUDA, &nvfp4_qpn4_prepare_sm70); @@ -323,6 +332,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("fp8_gemm_sm70_prefill_prescaled_out", torch::kCUDA, &fp8_gemm_sm70_prefill_prescaled_out); + // A distinct schema is also a capability marker: older extensions expose + // only the 8K-prefill contract and must not receive the new M=1 shapes. + ops.def( + "fp8_gemm_sm70_prescaled_m1_out(Tensor(a!) out, Tensor _in_feats, " + "Tensor _kernel, Tensor _prescaled_factors, int group_size, int k_ld, " + "int q_ld) -> ()"); + ops.impl("fp8_gemm_sm70_prescaled_m1_out", torch::kCUDA, + &fp8_gemm_sm70_prescaled_m1_out); + ops.def( "nvfp4_qpn2_prepare_sm70(Tensor weight_packed, Tensor weight_scale) -> " "Tensor[]"); @@ -402,6 +420,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int k_ld, bool gated_silu) -> ()"); ops.impl("sm70_f16_gemm_out", torch::kCUDA, &sm70_f16_gemm_out); + ops.def( + "sm70_glm_mhc_pre_norm_out(" + "Tensor gemm_mul, Tensor gemm_sqrsum, Tensor hc_scale, Tensor hc_base, " + "Tensor residual, Tensor(a!) post_mix, Tensor(b!) comb_mix, " + "Tensor(c!) layer_input, Tensor norm_weight, float rms_eps, " + "float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult, " + "int sinkhorn_repeat, float norm_eps) -> ()"); + ops.impl("sm70_glm_mhc_pre_norm_out", torch::kCUDA, + &sm70_glm_mhc_pre_norm_out); + + ops.def( + "sm70_glm_kda_fg_b_out(Tensor(a!) f_out, Tensor(b!) g_out, " + "Tensor f_input, Tensor g_input, Tensor f_weight, Tensor g_weight) -> " + "()"); + ops.impl("sm70_glm_kda_fg_b_out", torch::kCUDA, &sm70_glm_kda_fg_b_out); + ops.def( "sm70_f16_indexed_rerank_out(Tensor(a!) out, Tensor _in_feats, " "Tensor _kernel, Tensor candidate_ids, Tensor(b!) selected_raw, " @@ -635,6 +669,20 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("mxfp4_moe_dense_stage_sm70_out", torch::kCUDA, &mxfp4_moe_dense_stage_sm70_out); + ops.def( + "mxfp4_moe_qpn_m1_sm70_out(" + "Tensor(a!) out, Tensor input, Tensor weights, Tensor scales, " + "Tensor expert_ids, bool broadcast_input) -> ()"); + ops.impl("mxfp4_moe_qpn_m1_sm70_out", torch::kCUDA, + &mxfp4_moe_qpn_m1_sm70_out); + + ops.def( + "nvfp4_moe_qpn_m1_sm70_out(" + "Tensor(a!) out, Tensor input, Tensor weights, Tensor scales, " + "Tensor expert_ids, bool broadcast_input, int split_k) -> ()"); + ops.impl("nvfp4_moe_qpn_m1_sm70_out", torch::kCUDA, + &nvfp4_moe_qpn_m1_sm70_out); + ops.def( "nvfp4_moe_dense_stage_sm70_out(" "Tensor(a!) out, Tensor input, Tensor expert_offsets, " diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 0ae35468fa..d5d18f4ef1 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -133,16 +133,6 @@ Priority is **1 = highest** (tried first). | 7 | `FLASHINFER_MLA_SPARSE`**\*** | | 8 | `FLASHMLA_SPARSE` | -**Ampere/Hopper (SM 8.x-9.x):** - -| Priority | Backend | -| -------- | ------- | -| 1 | `FLASH_ATTN_MLA` | -| 2 | `FLASHMLA` | -| 3 | `FLASHINFER_MLA` | -| 4 | `TRITON_MLA` | -| 5 | `FLASHMLA_SPARSE` | - > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise. > > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details. @@ -224,6 +214,7 @@ MLA decode backends are selected using the standard | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | +| `GLM5_SM70_SPARSE` | fp16 | `auto`, `float16`, `fp8`, `fp8_e4m3` | Any | 512 | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 7.x | | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | diff --git a/docs/design/qwen36_27b_awq_mtp_target_verifier_fastpaths.md b/docs/design/qwen36_27b_awq_mtp_target_verifier_fastpaths.md index 6a62b80ad0..86facb379b 100644 --- a/docs/design/qwen36_27b_awq_mtp_target_verifier_fastpaths.md +++ b/docs/design/qwen36_27b_awq_mtp_target_verifier_fastpaths.md @@ -351,11 +351,11 @@ Required experiment: Primary source paths: -- https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/kernelDispatcher.h -- https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/kernelLauncher.h -- https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/kernel.h -- https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/utility.h -- https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp +- +- +- +- +- ### P1: SM90 Machete-style mixed-input WGMMA @@ -401,7 +401,7 @@ must profile 1-SM, 2-SM, Stream-K, and exact-M GEMV tactics for each of the four real shapes. Source: -https://github.com/NVIDIA/cutlass/blob/main/examples/86_blackwell_mixed_dtype_gemm/86_blackwell_mixed_dtype.cu + **Native NVFP4/MXFP4.** Blackwell's `tcgen05.mma.blockscaled` consumes native block-scaled FP4 and accumulates through TMEM. CUTLASS documents 2x instruction @@ -415,9 +415,9 @@ valid conversion. Sources: -- https://github.com/NVIDIA/cutlass/blob/main/examples/72_blackwell_narrow_precision_gemm/72a_blackwell_nvfp4_bf16_gemm.cu -- https://github.com/NVIDIA/cutlass/blob/main/examples/91_fp4_gemv/91_fp4_gemv.cu -- https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/llm_ptq/README.md +- +- +- ### TurboMind SM90 ceiling @@ -429,7 +429,7 @@ mainloop. Therefore, "enable TurboMind SM90" and "exploit Hopper's full mixed input path" are separate milestones. Source: -https://github.com/InternLM/lmdeploy/blob/main/src/turbomind/kernels/gemm/kernel/sm90_16816_4.cu + ### Ranked benchmark order @@ -538,11 +538,11 @@ or H100 target-forward baseline. ## Primary References -- NVIDIA Ampere tuning guide: https://docs.nvidia.com/cuda/ampere-tuning-guide/index.html -- NVIDIA Hopper tuning guide: https://docs.nvidia.com/cuda/archive/12.8.0/hopper-tuning-guide/index.html -- CUTLASS changelog and mixed-input GEMM support: https://github.com/NVIDIA/cutlass/blob/main/CHANGELOG.md -- Upstream TurboMind architecture-specific GEMM registry: https://github.com/InternLM/lmdeploy/tree/main/src/turbomind/kernels/gemm/kernel -- FlashInfer GDN decode/MTP API: https://docs.flashinfer.ai/api/gdn_decode.html -- FlashInfer GDN SM90 implementation contract: https://github.com/flashinfer-ai/flashinfer/blob/main/flashinfer/gdn_decode.py -- FlashInfer/TensorRT-LLM MTP attention API: https://docs.flashinfer.ai/generated/flashinfer.decode.trtllm_batch_decode_with_kv_cache.html -- TensorRT-LLM quantization hardware matrix: https://nvidia.github.io/TensorRT-LLM/latest/features/quantization.html +- NVIDIA Ampere tuning guide: +- NVIDIA Hopper tuning guide: +- CUTLASS changelog and mixed-input GEMM support: +- Upstream TurboMind architecture-specific GEMM registry: +- FlashInfer GDN decode/MTP API: +- FlashInfer GDN SM90 implementation contract: +- FlashInfer/TensorRT-LLM MTP attention API: +- TensorRT-LLM quantization hardware matrix: diff --git a/docs/design/sm70_dflash2_long_verify_decay.md b/docs/design/sm70_dflash2_long_verify_decay.md index 7a8f161214..8c92ac129d 100644 --- a/docs/design/sm70_dflash2_long_verify_decay.md +++ b/docs/design/sm70_dflash2_long_verify_decay.md @@ -448,11 +448,11 @@ respectively; 103/128 and 24/32 requests terminate naturally before the cap. The 32K/128K trace attribution is closed: target grouped attention is the long-context verification bottleneck. Pack-GQA48 is the retained candidate; its 1K/32K/128K full-graph endpoints, 1K-256K operator sweep, task scores, PPL, -and dataset-level acceptance gates are recorded above. The route remains -explicitly opt-in through `VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY`, so this -change does not silently alter other attention or speculation paths. Two -bounded follow-ups (tail compaction and page-address broadcast) are rejected -and must not be resurrected without new evidence. +and dataset-level acceptance gates are recorded above. At the initial public +merge, the route remained explicitly opt-in through +`VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY`. Two bounded follow-ups (tail +compaction and page-address broadcast) are rejected and must not be resurrected +without new evidence. The final source audit rebased all five patches without semantic change onto `main@acc0f6fb92`, after the generic prefix-anchored and E4M3 batch-XQA merges. @@ -463,3 +463,237 @@ alignment attribute preserves the packed verifier layout while restoring all or missing. All thirteen grouped-verifier cases pass on a V100 against this current-main build. No new full-model or 256K production-configuration claim is made by the current-main source audit. + +## 2026-08-26 current-main long-context follow-up + +This follow-up starts from `main@b06f54cee3` and keeps the practical production +contract fixed: Qwen3.8-27B NVFP4 target, E5M2 target KV, FP16 draft KV, TP4, +seven probabilistic draft tokens, prefix cache enabled, 256K maximum context, +4,096 maximum batched tokens, and FULL CUDA Graph. The frozen unprofiled +complete-round baseline is 18.477/21.583/30.641/42.300 ms at +1K/32K/128K/256K. The corresponding Nsight diagnostic values are +21.323/24.430/33.400/45.307 ms; the approximately 2.8-3.0 ms tracing overhead +is kept separate from service-wall results. + +The trace attribution remains consistent with the earlier audit. Sixteen target +full-attention layers dominate the context-dependent slope, while five draft +sliding-window attention layers contribute a smaller avoidable tail. The target +grouped partial kernel is launched once per full-attention layer with grid +`(1, 80)`, 512 threads, 128 registers per thread, and 50,944 bytes of dynamic +shared memory. Its retained single-layer baseline is approximately +0.048/0.206/0.746/1.451 ms at 1K/32K/128K/256K. + +### Draft sliding-window history skip + +The non-anchored draft paged-prefill kernel previously started from key tile +zero even though a 2,048-token sliding window masked all older tiles. Clamping +the first visited tile to `min_key_pos / BLOCK_N` removes that masked history; +the anchored/prefix path is deliberately unchanged. Exact DFlash2 shape +microbenchmarks report the following five-layer totals: + +| Context | Baseline | Candidate | Delta | +| --- | ---: | ---: | ---: | +| 1K | 0.737280 ms | 0.757760 ms | +0.020480 ms | +| 32K | 1.367040 ms | 1.264640 ms | -0.102400 ms | +| 128K | 1.730560 ms | 1.218560 ms | -0.512000 ms | +| 256K | 2.114560 ms | 1.233920 ms | -0.880640 ms | + +Baseline and candidate tensors are bitwise equal at all four lengths. The +targeted non-causal DFlash2 SWA and anchored-SWA suites pass all 21 cases. The +small 1K operator movement is 0.02048 ms per complete round and remains subject +to the full-model one-percent promotion gate. + +### Volta shared-row bank-conflict removal + +The packed grouped verifier used power-of-two shared row strides for Q/K/V +(`256` halves), scores (`32` floats), and probabilities (`32` halves). Those +strides repeatedly present the same bank pattern to Volta WMMA shared loads. +The initial candidate padded Q/K/V rows by eight halves and both score and +probability rows by eight elements. This changes only shared-memory addressing: + +- Q and K/V row stride: `256 -> 264`; +- score and probability row stride: `32 -> 40`; +- softmax arithmetic, FP32 reductions, visibility masks, split count, launch + geometry, output workspace, and sampling remain unchanged. + +The candidate uses 54,528 bytes of opt-in dynamic shared memory and retains one +CTA per SM and 128 registers per thread. Q-only and K/V-only ablations improved +the 32K/128K/256K one-pass kernel to +0.168/0.592/1.147 ms and 0.184/0.618/1.200 ms, respectively. Padding both sides +reached 0.136/0.469/0.900 ms; padding the score/probability rows as well produced +the retained result below. + +The production-page-size CUDA-Graph A/B/A medians use 40 warmups and 200 timed +rounds on the same V100: + +| Context | Baseline A1 | Padded candidate | Baseline A2 | Delta vs. A mean | +| --- | ---: | ---: | ---: | ---: | +| 1K | 0.051200 ms | 0.031744 ms | 0.044032 ms | -33.33% | +| 32K | 0.206848 ms | 0.124928 ms | 0.204800 ms | -39.30% | +| 128K | 0.745472 ms | 0.419840 ms | 0.746496 ms | -43.72% | +| 256K | 1.451008 ms | 0.807936 ms | 1.451008 ms | -44.32% | + +Candidate outputs are bitwise identical to the retained baseline at all four +lengths, and all thirteen grouped-verifier GPU cases pass. The raw A/B/A records +are `grouped-sp40-aba-a1.json`, `grouped-sp40-aba-b.json`, and +`grouped-sp40-aba-a2.json` under the task microbenchmark directory. + +Two bounded alternatives were rejected and fully reverted. A 192-thread, +160-split layout intended to admit two CTAs per SM slowed the four points to +0.063/0.279/0.889/1.703 ms. An in-CTA QK-next/PV-current software pipeline used +97,280 bytes of shared memory and slowed them to +0.058/0.264/0.972/1.900 ms. Both were numerically correct, but neither improved +the measured critical path; they must not be revived without new hardware +evidence. + +The combined candidate was then run with the frozen TP4 production contract. +The four-request unprofiled diagnostic improves steady decode from +147.5/259.0/176.6 to 170.3/317.7/240.0 token/s at 32K/128K/256K. The 1K request +is excluded from this throughput comparison because both runs include first-use +JIT. These rates remain acceptance-sensitive and are not used as the verifier +latency proof. + +Independent Nsight Systems traces provide that proof. Each result drops the +first graph round and averages all four ranks; tracing overhead is therefore +kept separate from unprofiled service wall. Values are milliseconds per +complete verification cycle: + +| Context | Phase | Baseline | Candidate | Delta | +| --- | --- | ---: | ---: | ---: | +| 1K | Draft | 3.774 | 3.933 | +0.159 | +| 1K | Draft to target | 0.251 | 0.297 | +0.046 | +| 1K | Target verify | 16.003 | 15.881 | -0.122 | +| 1K | Target to draft | 1.295 | 1.304 | +0.009 | +| 1K | Complete | 21.323 | 21.415 | +0.092 (+0.43%) | +| 32K | Draft | 4.248 | 4.215 | -0.033 | +| 32K | Draft to target | 0.255 | 0.285 | +0.030 | +| 32K | Target verify | 18.623 | 17.403 | -1.220 | +| 32K | Target to draft | 1.303 | 1.307 | +0.003 | +| 32K | Complete | 24.430 | 23.210 | -1.219 (-4.99%) | +| 128K | Draft | 4.598 | 4.342 | -0.256 | +| 128K | Draft to target | 0.256 | 0.332 | +0.076 | +| 128K | Target verify | 27.239 | 22.179 | -5.060 | +| 128K | Target to draft | 1.307 | 1.318 | +0.012 | +| 128K | Complete | 33.400 | 28.171 | -5.229 (-15.65%) | +| 256K | Draft | 5.091 | 4.443 | -0.648 | +| 256K | Draft to target | 0.254 | 0.256 | +0.002 | +| 256K | Target verify | 38.651 | 28.377 | -10.274 | +| 256K | Target to draft | 1.311 | 1.308 | -0.004 | +| 256K | Complete | 45.307 | 34.384 | -10.923 (-24.11%) | + +The complete-round medians are 21.216/23.125/27.910/34.090 ms, versus +21.320/24.435/33.399/45.219 ms for the baseline. The mean 32K-to-256K growth +falls from 20.877 to 11.173 ms (-46.48%), while the 1K point stays within the +one-percent non-regression gate. This closes the speed and long-context-decay +gates for the kernel change. + +The unprofiled four-request run has exactly the same aggregate mean acceptance +length (`4.932692`), accepted count (`409`), and drafted count (`728`) as the +frozen baseline. One 1K probabilistic trace differs by -0.058824 in mean +acceptance, narrowly outside the single-sample 0.05 boundary; the source +operator outputs are nevertheless bitwise equal at all four audited lengths. +Consequently no quality-promotion claim is made from the synthetic traces. The +frozen dataset-level acceptance, task-score, and WikiText PPL gates remain +required before default promotion. + +### PR 285 fixed-layout follow-up + +The operator reference identified for this follow-up is +[PR 285](https://github.com/1CatAI/1Cat-vLLM/pull/285), not PR 206. PR 285 +improves the independent batch-one E4M3 XQA route from 40.561 to 61.834 token/s +at 128K and from 27.456 to 50.376 token/s at 256K. Its relevant retained +mechanisms are a compile-time interleaved physical stride and split-local page +ID staging. The DFlash2 verifier already combines QK, online softmax, and WMMA +P@V in one 80-CTA kernel, so PR 285's merged long-wave launch is already the +local topology. Its scalar shared-V `half2` P@V loop is not copied because that +would replace the verifier's tensor-core P@V arithmetic rather than merely +specialize addressing. + +The new dispatch admits a fixed layout only for the exact production unbound +views from `[blocks, 2, page, 1, 256]`, page size 1,648 or 3,296, and matching +physical strides. All other layouts retain the generic stride path. The fixed +kernel removes runtime 64-bit block-stride products, and stages the few page +IDs used by each context split into shared memory. The page IDs share the +existing Q-publication barrier, so this adds no CTA barrier. Both changes have +independent default-on rollback flags: + +- `VLLM_FLASH_V100_DFLASH2_FIXED_INTERLEAVED=0` disables fixed addressing; +- `VLLM_FLASH_V100_DFLASH2_STAGE_PAGE_IDS=0` disables page-ID staging. + +The same audit found and fixed a generic paired-load defect: its physical +interleaved block stride was hard-coded to page 800. It is now +`2 * page_size * head_dim`. This protects page-1,568 E4M3 XQA while admitting +the new page-1,648/3,296 DFlash2 layouts. The existing page-1,568 E4M3 scalar +and wave-partition CUDA tests pass all fourteen parametrized cases after the +change. + +A same-process CUDA-Graph comparison on the production page-3,296 layout keeps +the initial score/probability-padding binary fixed and toggles only fixed +addressing plus page staging: + +| Context | Generic | Fixed plus staged | Delta | +| --- | ---: | ---: | ---: | +| 1K | 0.035840 ms | 0.035840 ms | 0.00% | +| 32K | 0.142336 ms | 0.140288 ms | -1.44% | +| 128K | 0.447488 ms | 0.440320 ms | -1.60% | +| 256K | 0.801792 ms | 0.789504 ms | -1.53% | + +Every output is bitwise equal. Separate ablations attribute approximately +0.7-1.3% to the fixed physical stride and a further 0.4-0.7% to page staging. +This is deliberately a small address-path improvement, not a claim that the +full PR 285 throughput ratio transfers to the different DFlash2 WMMA kernel. + +An A/B/A shared-row follow-up then removed score padding while retaining +probability padding. On the same physical GPU, the selected +Q/K/V-264, score-32, probability-40 layout compares with the earlier +Q/K/V-264, score-40, probability-40 layout as follows: + +| Context | Selected score-32 | Earlier score-40 | Delta | +| --- | ---: | ---: | ---: | +| 1K | 0.030720 ms | 0.030720 ms | 0.00% | +| 32K | 0.120832 ms | 0.121856 ms | -0.84% | +| 128K | 0.407552 ms | 0.410624 ms | -0.75% | +| 256K | 0.784384 ms | 0.789504 ms | -0.65% | + +The output SHA256 is identical across the two separately built extensions at +every length. The final extension SHA256 is +`b6952509ffb117c16ba43720c1affd5024f350ccd2ca4324d89022953692bc96`. +Its production fixed/staged one-pass specialization uses 127 registers per +thread, zero stack, and 52,992 bytes of dynamic shared memory; launch geometry +and one-CTA-per-SM occupancy are unchanged. All fifteen grouped-verifier GPU +cases, all fourteen page-1,568 E4M3 regression cases, and all 56 environment +tests pass. + +The bounded probabilistic mixed-quality A/B uses 32 shuffled requests, a +2,048-token cap, natural EOS, and the production sampling contract. The first +control and candidate both score 19/32, with identical per-suite pass counts: +GSM8K 7/9, HumanEval 7/9, MATH500 5/7, and the known-invalid embedded MBPP +subset 0/7. The embedded MBPP result is useful only as a relative regression +check; the valid independent MBPP quality gate is recorded earlier. Candidate +natural termination improves from 25/32 to 26/32. Request-mean acceptance +changes from 4.809806 to 4.841060 (+0.031254), while length-weighted pooled +acceptance changes from 4.226907 to 4.154068 (-0.072839). + +An identical-control replay proves that separate probabilistic engine starts +are not token-stream reproducible: only 4/32 streams match the first control, +despite no implementation or configuration change. Its task score is 18/32, +pooled acceptance is 4.190773, and request-mean acceptance is 4.834261. Against +this adjacent replay, the fixed/staged candidate is one task higher, changes +pooled acceptance by -0.036705, changes request-mean acceptance by +0.006799, +and stays inside the frozen 0.05 acceptance gate. This evidence rejects a +quality-regression attribution; it does not claim that address specialization +improves model quality. Operator outputs remain bitwise equal and the earlier +valid 128-request, independent-MBPP, and WikiText gates continue to define the +absolute production-quality envelope. + +### Default promotion + +The private current-main audit promotes the grouped verifier to default-on for +its existing exact runtime contract. Admission remains shape- and route-based: +SM70, DFlash selector verification, Q8/H6/Hkv1/D256, E5M2 KV, page 1,648 or +3,296, one request, causal full attention, and a model length of at least 32K. +It does not inspect a model or checkpoint identity. The retained operator is +bitwise equal to the prior route, the inherited dataset/PPL gates remain valid, +and the 1K complete-round movement stays inside the one-percent guard. Set +`VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY=0` to restore the independent-row +fallback. diff --git a/docs/design/sm70_dflash2_lookup_augmented.md b/docs/design/sm70_dflash2_lookup_augmented.md new file mode 100644 index 0000000000..c8409e3ea3 --- /dev/null +++ b/docs/design/sm70_dflash2_lookup_augmented.md @@ -0,0 +1,171 @@ +# SM70 DFlash2 lookup-augmented block drafting + +## Purpose and provenance + +This change ports lookup-augmented block drafting (LABD) to the public 1Cat +MRV2 DFlash2 path. The neural checkpoint remains a block-8 model (one anchor +plus seven trained draft positions), while the target may verify a q16 block +when the request is demonstrably reproducing its own context. Normal and +low-confidence traffic stays on q8. + +The algorithm is adapted from +[`syv-ai/qwen38-27b-rtx3090`](https://github.com/syv-ai/qwen38-27b-rtx3090) +revision `69ba4d0688c6ae76cb9d3c4a5c3b36445e1b040c`, specifically the +Apache-2.0 `patches/dflash2-lookup-drafting.patch`. The implementation is +rebased onto 1Cat's Model Runner V2, sparse probabilistic rejection sampler, +SM70 selector, hybrid KV manager, and Flash-V100 verifier rather than applying +the patch mechanically. Development is tracked in public Draft PR #355. + +## Runtime contract + +- `method=dflash`, `ngram_assist=true`, and `num_speculative_tokens=15` + activate LABD for a selector-capable DFlash2 checkpoint whose trained block + is smaller than the configured verifier. +- The DFlash2 model, grouped convolution, selector lattice, and draft KV retain + the checkpoint-native seven draft positions. Lookup owns at most the eight + additional target positions. +- One Triton program per active request scans the authoritative UVA int32 token + history. It selects the longest suffix match and breaks ties by recency; + overlapping matches are legal. +- The lookup proposal is fused with the neural proposal. Filled probabilistic + rows become point masses in the existing sparse draft-logit cache, including + complete erase metadata for the following step. +- Structured-output and prefill batches retain q8 and do not use lookup. +- The host controller enters q16 only after two consecutive strong copy + signals. B1 may coast for three steps; batches larger than one never keep + sticky state across a miss. +- Both q8 and q16 target graphs are captured. The DFlash draft graph remains + q8 because that is the checkpoint contract. + +The asynchronous scheduler cannot carry a worker-selected proposal count back +to the scheduler: it pads every step to the configured maximum. Adaptive LABD +therefore disables asynchronous scheduling by default. An explicit async +request fails at startup instead of silently becoming always-q16. Setting +`VLLM_DFLASH2_LOOKUP_ADAPTIVE=0` deliberately selects fixed-width q16 and may +retain async scheduling. + +## Cache and Flash-V100 integration + +The wider verifier changes the Mamba-align hybrid page from 1,648/3,296 to +1,728/3,456 elements. The hybrid KV reservation uses the target verification +width while all DFlash model allocations use the trained width. Flash-V100's +grouped verifier accepts all four page layouts. The existing CUDA extension's +runtime-stride fallback handles 1,728/3,456 exactly; a fixed specialization is +only justified if the paired microbenchmark shows a material gap. + +## Evidence, 2026-08-28 + +CPU configuration, routing, controller, graph-layout, and policy tests pass. +The controller test proves the sequence q8, q8, q16 for two consecutive strong +signals and proves that a multi-request miss cannot coast on q16. The config +tests prove default async disablement and a targeted error for explicitly +enabled async scheduling. + +On Tesla V100-SXM2-32GB, the new 1,728/3,456 page layouts match an FP32 +reference for both q8 and q16 (four strict cases). The six-case GPU LABD suite +passes suffix hit/miss/overlap, request eligibility, 7-to-15 fusion, controller +signal boundaries, and sparse point-mass cache invariants. + +The standalone lookup-plus-fusion CUDA Graph cost is: + +| Context | B1 | B2 | B4 | B8 | +| ---: | ---: | ---: | ---: | ---: | +| 1K | 0.0111 ms | 0.0131 ms | 0.0139 ms | 0.0152 ms | +| 32K | 0.1210 ms | 0.1254 ms | 0.1274 ms | 0.1351 ms | +| 64K | 0.2336 ms | 0.2359 ms | 0.2317 ms | 0.2357 ms | +| 128K | 0.4182 ms | 0.4156 ms | 0.4124 ms | 0.4705 ms | + +The external benchmark artifact is `lookup-uva-sm70.json`. The nearly flat +B1-to-B8 slope confirms that the lookup kernel exposes request parallelism +rather than serializing the batch. + +The runtime-stride 3,456-page grouped-verifier CUDA Graph measurements are: + +| Query/context | Grouped | Independent XQA | Speedup | +| --- | ---: | ---: | ---: | +| q8 / 1K | 0.0320 ms | 0.0564 ms | 1.76x | +| q8 / 32K | 0.1240 ms | 0.9098 ms | 7.34x | +| q8 / 128K | 0.4095 ms | 2.9715 ms | 7.26x | +| q16 / 1K | 0.0273 ms | 0.0479 ms | 1.76x | +| q16 / 32K | 0.2167 ms | 1.5883 ms | 7.33x | +| q16 / 128K | 0.7855 ms | 5.6914 ms | 7.25x | + +The 32K/128K movement versus the old fixed 3,296 page tracks the 4.85% larger +page closely. A new fixed 3,456 specialization would recover only the earlier +address-path margin (roughly 1--2%), so it is deferred until a full-model trace +shows that margin on the critical path. + +An initial q15 route smoke produced coherent lookup proposals and positions +beyond the neural seven-token block, but it ran with async scheduling. Its +`75` drafts over five rounds are proof that async pinned every round to 15, not +valid evidence for the adaptive policy or throughput. That run is retained as +a route/correctness diagnostic only. + +The first synchronous end-to-end pair uses the production NVFP4 target, TP4, +E5M2 target KV, FP16 draft KV, prefix caching, Mamba align, a 4,096-token +prefill budget, probabilistic sampling (`temperature=1`, `top_p=0.95`, +`top_k=20`), and CUDA Graphs. Each ordinary-path number below is the second +256-token request in one engine process, after inference-time Triton JIT has +already completed: + +| Workload | Verifier behavior | Mean acceptance length | Steady decode | TPOT | +| --- | --- | ---: | ---: | ---: | +| ordinary async rate-limiter task, q7 baseline | fixed q8 | 3.048 | 168.52 tok/s | 5.934 ms | +| same task, adaptive LABD | q8 throughout | 3.122 | 170.98 tok/s | 5.849 ms | +| repeated-context continuation, adaptive LABD | q8 then sustained q16 | mixed-run 4.896 | 316.27 tok/s | 3.162 ms | + +The normal sample is 1.46% faster than the paired q7 sample, so this route did +not expose a normal-request throughput penalty. It is only one prompt and is +not a promotion-level confidence interval. The mixed normal/copy counters +contain 106 measured rounds and 862 proposals, which implies 91 q8 rounds and +15 q16 rounds. Positions 8--15 were accepted in 14 of those q16 rounds. The +controller log records the normal request as ineligible and changes to q16 +only after the copy request produces the required consecutive strong signals. + +The paired runs use probabilistic sampling and their target trajectories split +at token 73. They are therefore performance evidence, not an exact-token +quality claim. Dataset scores and parser validity remain the quality oracle. + +The fixed-sample API quality battery gives the same scores as target-only/q7: + +| Gate | Adaptive LABD | Target-only/q7 reference | +| --- | ---: | ---: | +| BFCL-style function calling | 29/32 | 29/32 | +| strict JSON Schema | 7/8 | 7/8 | +| ToolACE subset | 12/12 | 12/12 | +| NexusRaven subset | 13/16 | 13/16 | +| structured B1 | 12/12 | 12/12 | +| structured B4 | 12/12 | 12/12 | +| prefix-state/cache checks | 5/5 | 5/5 | + +One BFCL parallel-call case initially appeared to regress to 28/32: the final +`math.factorial(number=15)` arrived as `{}`. Target-only and LABD emitted +identical token IDs, proving that this was not a model or acceptance failure. +The q15 burst placed a complete function body in the parser buffer before its +header delta was emitted, and the next/final delta contained only the closing +tag. The Qwen3-Coder streaming parser now consumes an already-complete body +when it emits that header. Its exact 13-burst reproducer and all 104 parser +tests pass; BFCL returns to 29/32. + +A separate cold-cache run exposed the same MRV2 coverage gap in startup +warmup: the old model-module scan did not see the GDN cache bound in the static +forward context. A strict JSON request could therefore compile the ordinary +causal-conv update during inference and took 218.45 seconds. MRV2 now warms one +bound Qwen GDN layer with the production cache layout before graph capture. On +a new Triton and Inductor cache, the startup log confirms that warmup, the +first strict JSON request returns correct output in 4.14 seconds, and +inference-time `_causal_conv1d_update_kernel` JIT falls from one to zero. Other +first-request auxiliary kernels still account for most of those four seconds +and are tracked separately from decode performance. + +## Remaining promotion gates + +1. Compare q7 DFlash2, adaptive LABD, and fixed q16 at B1/B2/B4. Report target + verify time, full round time, tokens per round, per-stream decode, aggregate + decode, TTFT, resident requests, and preemptions. +2. Extend the now-passing protocol battery to broad coding and long-context + suites, then sweep all 128 relevant CUDA-Graph/prefix-cache residues. No + task-score or parser-validity regression is acceptable. + +LABD remains opt-in until every gate above is complete. Drafter-free chaining +is a separate second-stage experiment and stays disabled in this change. diff --git a/docs/design/sm70_dflash2_prefill_closure.md b/docs/design/sm70_dflash2_prefill_closure.md new file mode 100644 index 0000000000..1621b4c077 --- /dev/null +++ b/docs/design/sm70_dflash2_prefill_closure.md @@ -0,0 +1,171 @@ +# SM70 DFlash2 Prefill Closure + +## Scope and integration base + +This private campaign is stacked on the quality-audited DFlash2 branch at +`ee4ac48a479c3dbd458d5f7c09a59f39fd271d82`. It keeps the accepted NVFP4 +target, official BF16 DFlash2 draft, FP8 E5M2 target KV, FP16 draft KV, +prefix caching, Mamba alignment, and CUDA Graph decode contract. Draft-MLP +QPN8 remains disabled. + +The first objective is to restore already accepted SM70 prefill operators in +source-overlay deployments. Kernel arithmetic is not changed by that repair. +Any later shape expansion is a separate gate. + +## Project PR audit + +The historical short- and long-prefill numbers use different contracts: + +| Evidence | Contract | Accepted result | Boundary | +|---|---|---:|---| +| Public PR #271 | Qwen3.8-27B-FP8, TP4, exact input 8000, target-only | 5121.44 request-wall and 5170.96 pure-prefill tok/s | Exact-8K only | +| Public PR #324 | Same exact-8K FP8 contract | 5500 tok/s is a campaign target, not a measured implementation result | Documentation-only PR | +| Public PR #224 | Qwen3.8-27B-FP8, TP4, exact input 65536, target-only | 2798.6 to 3496.4 prompt tok/s after the exact D256 operator | Closest retained 64K target-only reference | +| Private PR #8/#13 | Qwen3.8-27B-FP8, TP4, input 261888, chunk 8192 with FP16 Mamba/SSM cache and Q8000 aligned chunks, target-only | 2438.89 prompt tok/s | Stable max-aware D256 architecture | +| Rejected public PR #315 lane | Same 256K FP8 contract | 2971.51 prompt tok/s | Rejected: 32 output token IDs were zero | + +The 2438.89 tok/s route uses max-shifted exponentiation and max-aware online +softmax merging. Its output hash exactly matched the exact control. It also +explicitly overrides the checkpoint's FP32 SSM-cache contract to FP16; that +override is a separate quality variable and is not inherited automatically by +DFlash2. The removed raw-logit half2 polynomial must not be restored. + +## Current DFlash2 baseline + +The same-card cold benchmark resets prefix cache before every warmup and +measurement. It uses the practical 256K API contract, including chunk 4096. + +| Input | Mean pure prefill | Prompt throughput | +|---:|---:|---:| +| 32768 | 10.485347 s | 3125.12 tok/s | +| 65536 | 25.317796 s | 2588.54 tok/s | + +All three repeats at each length emitted the same first-token hash. Artifact: +`/data/minimax-h3/task-cache/v100-dflash2-prefill-32k64k-20260827/current-dflash2-cold-prefill-v1/`. + +## Confirmed root cause + +Every retained practical DFlash2 long-prefill log reports that +`_vllm_fa2_C` cannot be imported. The source checkout contains the D256 +dispatch and quality-safe architecture, but the source overlay shadows the +installed package containing the native extension. Therefore the merged D256 +path is not merely underperforming; it has never executed in these runs. + +The accepted stable binary is retained at +`/data/minimax-h3/task-cache/qwen38-d256-attn-80tflops-20260825/build/exact-stat-256k-py312-v2/_vllm_fa2_C.abi3.so` +with SHA256 `f9f9acbc610c87fce9984e8fbd93fe0c8fa59887542123a74b3eaef6d3b8abf9`. +It loads against the active Torch 2.10/CUDA 12.8 environment and registers the +required dense, paged, split-KV3, and stable GQA architecture operators. + +This branch adds an explicit `VLLM_SM70_FA2_D256_LIBRARY` source-overlay +sidecar. It is opt-in and follows the existing SM70 native-sidecar convention. +Bundled-wheel behavior remains unchanged, and missing or incompatible +operators still fail closed to the existing fallback with a warning. +Both the benchmark preflight and runtime loader validate registered operators, +not merely a successful Python-interface import. This covers partially cached +interfaces that otherwise appear importable while exposing no native kernels. + +## Shape boundary and next measurements + +The practical chunk-4096 contract can use the exact D256 Split-D operators, +but it cannot enter the stable long GQA architecture, whose validated kernel +contract is Q8000 with KV16K..256K in 8K steps. With seven speculative slots, +the scheduler first reduces the configured 4096-token budget to 4089. The +checkpoint's FP32 SSM state makes one aligned attention/Mamba block 1648 +tokens, so the observed steady prefill query is Q3296. A configured 8192-token +budget retains that FP32 state and yields Q6592, not Q8000. The historical +target-only FP16 Mamba/SSM contract used an 800-token block, but DFlash2 grows +the convolution state by seven speculative slots. Its FP16 block is therefore +880 tokens and the same budget yields Q7920, so it cannot enter the existing +Q8000 architecture either. The next paired measurements are therefore +deliberately separated: + +1. chunk 4096 plus the stable sidecar, to measure the dependency-closure gain; +2. chunk 8192 plus the same sidecar and FP32 SSM state, to measure Q6592; +3. chunk 8192 plus FP16 Mamba/SSM cache, to measure the actual Q7920 DFlash2 + geometry rather than crediting the target-only Q8000 route; +4. profile the remaining NVFP4 projection cost before considering a new + Q6592/Q7920 attention architecture. + +Each candidate must prove operator-route hits, preserve output validity, fit +the 256K DFlash2 memory contract, and retain the quality-audit PPL and scored +coding gates. Prefix-hit time is reported separately and never counted as cold +prefill throughput. + +## Dependency-closure A/B + +The first candidate changes only native-extension resolution and keeps chunk +4096. All six measured requests are cold. The stable sidecar loads on all four +ranks and the benchmark reports both required D256 operators as available. + +| Input | Missing-extension control | Stable-sidecar candidate | Throughput gain | +|---:|---:|---:|---:| +| 32768 | 3125.12 tok/s | 3476.53 tok/s | +11.24% | +| 65536 | 2588.54 tok/s | 3103.02 tok/s | +19.87% | + +Candidate pure-prefill means are 9.425490 s and 21.120102 s. The three repeats +at each length retain the control first-token hash +`54363ddee68f4a5db81c9d37e5fb738d28f5b67dc7f725ad7333172b1ea157da`. +Artifact: +`/data/minimax-h3/task-cache/v100-dflash2-prefill-32k64k-20260827/candidate-stable-fa2-q4096-v1/`. + +## Chunk and Mamba dtype closure + +Increasing the configured budget to 8192 while retaining the checkpoint's +FP32 SSM state produces Q6592. It is neutral at 32K and slightly slower at 64K: + +| Input | Q3296 sidecar | Q6592 sidecar | Change | +|---:|---:|---:|---:| +| 32768 | 3476.53 tok/s | 3481.83 tok/s | +0.15% | +| 65536 | 3103.02 tok/s | 3077.03 tok/s | -0.84% | + +The FP16 Mamba/SSM arm produces Q7920, not Q8000, and loses at both lengths: + +| Input | Q3296 sidecar | Q7920 FP16 state | Change | +|---:|---:|---:|---:| +| 32768 | 3476.53 tok/s | 3271.49 tok/s | -5.90% | +| 65536 | 3103.02 tok/s | 2827.03 tok/s | -8.90% | + +All twelve measured requests retain the same first-token hash. Q6592 is not a +promotion, and Q7920 is rejected on speed before spending a dataset-quality +run. Artifacts are `candidate-stable-fa2-q8192-fp32-v4` and +`candidate-stable-fa2-q8000-mamba-fp16-v1` under the task root above. + +## NVFP4 projection candidate + +The accepted DFlash2 target keeps M<=8 verification on the existing QPN2 +route, while large-M prefill currently falls through to TurboMind W4A16. A +single-V100 M3296 microbenchmark shows that the already present bounded FP16 +QPN4 prefill operator is materially faster even though its timing includes +weight dequantization on every call: + +| Projection | TurboMind | bounded QPN4 | Latency reduction | +|---|---:|---:|---:| +| fused gate/up, 5120x8704 | 5.351 ms | 3.896 ms | 27.19% | +| down, 4352x5120 | 2.646 ms | 1.866 ms | 29.47% | + +The candidate reuses the QPN2 code and E4M3 scale-code buffers that are already +resident for verification, plus the existing shared 85-MiB FP16 workspace. It +does not retain a third packed weight layout: QPN2 and QPN4 expose different +2-D shapes but use the same flattened physical tile order, so the bridge uses +zero-copy views. Admission is default-off and requires +`VLLM_SM70_NVFP4_QPN2_PREFILL=1`; M below the separately recorded crossover +threshold keeps TurboMind, and M<=8 verification remains QPN2. The current +campaign stops at 64K; no 128K/256K throughput run is required for promotion. +The single-V100 bridge probe confirms byte-identical flattened code and scale +storage despite shapes `[N,K/2]` and `[K,N/2]`; both ordinary and fused +gate-SiLU QPN4-prefill calls are bitwise equal with maximum absolute error +zero. + +## Rejected SWA tail-only shortcut + +The DFlash2 draft uses a 2048-token sliding window, but projecting and writing +only the currently readable tail is not valid with automatic prefix caching. +Draft-group blocks are registered by content for later requests; a block whose +target tokens were computed without the matching draft-context KV write can be +reused later at a sequence length where those tokens are inside the window. +That produces a cache hit backed by uninitialized draft KV and collapses +acceptance. Therefore this campaign keeps full coverage of every newly +computed target token. A future tail-only design would first need to decouple +draft-group cache registration from target-prefix hits and prove rebuild +semantics; that is outside this prefill dependency closure. diff --git a/docs/design/sm70_dflash2_quality_audit.md b/docs/design/sm70_dflash2_quality_audit.md new file mode 100644 index 0000000000..d0c6371db8 --- /dev/null +++ b/docs/design/sm70_dflash2_quality_audit.md @@ -0,0 +1,316 @@ +# SM70 NVFP4 DFlash2 output-quality audit + +## Scope and provenance + +This audit starts from private integration commit +`d63e9490f65f9e01f6649053c1ab72922034b931` on branch +`codex/v100-dflash2-quality-audit-20260826-124851`. It covers the practical +Qwen3.8-27B NVFP4 target plus official BF16 DFlash2 drafter on TP4 V100. Remote +deployment and packaging are out of scope. + +The production contract remains temperature 1.0, top-p 0.95, top-k 20, +`xhigh` reasoning, seven probabilistic draft tokens, FP8 E5M2 target KV, FP16 +draft KV, FULL target/draft CUDA Graphs, prefix caching, Mamba align, 256K +maximum context, 4096 maximum batched tokens, and the Qwen tool/reasoning +parsers. Quality is judged by scored distributions rather than greedy or +bitwise identity. + +Upstream was checked once on 2026-08-26. vLLM PR +[`#52816`](https://github.com/vllm-project/vllm/pull/52816) is now merged. Its +following draft-logit cache-stride fix, vLLM PR `#53017`, does not apply to +this DFlash2 path: DFlash2 uses its own sparse top-16 cache kernel and passes +both request and step strides explicitly. Open quantized-drafter fixes are +also outside this contract because the retained drafter is the official BF16 +checkpoint. + +The structured-output audit found a separate relevant upstream chain after +the local May 29 fork point: vLLM PRs +[`#44297`](https://github.com/vllm-project/vllm/pull/44297), +[`#44993`](https://github.com/vllm-project/vllm/pull/44993), +[`#52805`](https://github.com/vllm-project/vllm/pull/52805), and +[`#53046`](https://github.com/vllm-project/vllm/pull/53046). They constrain +draft rows after a reasoning boundary, advance only the post-boundary suffix, +stop XGrammar at termination, and avoid logging expected invalid pre-mask +drafts as FSM errors. This is directly relevant to the still-open upstream +DFlash2 `json_object` report +[`#53777`](https://github.com/vllm-project/vllm/issues/53777); only this minimal +behavioral closure is backported here. + +## Existing evidence + +- The repaired stable RMSNorm extension SHA `7689e5...d449` passes its 13-case + batched-weight suite. The stale SHA `fe986a...9ba8c` reproduced the historical + row-zero context-K normalization bug and is not valid quality evidence. +- Draft-MLP QPN8 was removed after request-mean acceptance changed by + `-0.06397`, beyond the fixed `-0.05` gate. It remains disabled. +- The corrected MBPP-32 16K pair scores target-only/DFlash2 `31/32` versus + `32/32`; EvalPlus base `30/31` versus `31/31`; and EvalPlus plus `27/31` + versus `30/31`. HumanEval-32 EvalPlus is `32/32` for both. This is useful but + is only one sampling realization. +- The practical selector-QPN8 run on 64 GSM8K rows changed request-mean + acceptance `5.39851 -> 5.43951`, but its speed comparison used different + physical GPU groups. The design record explicitly kept the route opt-in. +- The packaged practical launch nevertheless forced + `VLLM_SM70_DFLASH2_QPN8_RERANK=1` and + `VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER=0`. This bypasses the source defaults + (`rerank=0`, `dense_order=1`) and must not be treated as a promoted quality + contract. +- Historical real-hidden candidate-order shadows contain no missing dense + top-20 token from the QPN8 top-64 support in the measured sample, and their + returned FP16 values match. They still change the local top-20 token set on + 69/4640 top-20 rows (1.49%) and top-1 on 6/4640 rows because equal-valued + cutoff tokens are resolved in candidate order. +- An offline scan of 2,048 retained selector-alignment dumps covers 16,384 + target rows. The 19th and 20th returned FP16 logits are equal on 544 rows + (3.3203%); four rows expose three equal cutoff entries within the returned + support. Visible cutoff probability mass averages about `0.000445`. This is + small per step but can accumulate in multi-thousand-token coding outputs. +- The retained 32K eager LongBench pair has four Qasper/NarrativeQA records. + Target-only and DFlash2 have identical token hashes on all four records and + the same average score, `36.6667`. This is useful long-prompt parity evidence, + but it does not replace Graph, prefix-cache, or 128K/256K validation. +- The retained grouped verifier is lower risk: its dedicated operator and + graph-replay tests are bitwise equal to the accepted fallback under the + admitted SM70 shape. Fixed interleaved page addressing and staged page IDs + still require the paired state-pollution checks below. + +## Hypotheses, in test order + +1. **Experimental selector launch configuration.** Candidate-order QPN8 may + alter low-probability target support at FP16 ties or very rarely miss the + dense top-20 support. Dense-order QPN8 should retain most of the speed while + restoring the local dense-vocabulary tie contract. +2. **Compact TP top-k boundary.** Sparse target rejection merges per-rank + candidates rather than materializing the full vocabulary. Non-tied rows are + exact; tied cutoff rows need a real-hidden global shadow before the path is + called distribution-equivalent to target-only. +3. **Graph/prefix/request-slot state.** Same-seed requests must be checked cold, + warm, after mixed-length contamination, and after prefix-cache reuse. A + target-only repeat is required before attributing cross-process stochastic + variation to DFlash2. +4. **Sparse rejection arithmetic.** Temperature application, target top-k then + top-p masking, conditional selector scores, token-keyed acceptance uniforms, + and `log1p` residual sampling are statically consistent. Dense versus sparse + same-source scoring remains the decisive dynamic check. + +## Minimal isolation matrix + +Every arm uses the same source, extension hashes, physical GPU group, model, +prompt order, request seeds, natural EOS, and practical engine parameters. + +| Arm | DFlash2 | Sparse target rejection | Selector QPN8 | Dense order | Purpose | +| --- | --- | --- | --- | --- | --- | +| `T0` | no | n/a | n/a | n/a | target-only score and repeat variance | +| `D0` | yes | off | off | n/a | dense DFlash2 reference | +| `D1` | yes | on | off | n/a | isolate compact rejection | +| `D2` | yes | on | on | on | quality-first accelerated selector | +| `D3` | yes | on | on | off | reproduce current packaged launch | + +The first screen uses a small fixed coding/reasoning subset and at least three +predeclared seeds. It records token/text hashes as diagnostics but promotes by +paired scores, natural-stop/invalid rates, acceptance, and confidence +intervals. Eager real-hidden shadows are run only long enough to measure QPN8 +support coverage and compact-global top-k boundary behavior. Graph runs then +check immediate repeats, mixed-length contamination, cold/warm prefix cache, +and request-slot reuse. + +The full gate uses HumanEval and MBPP with EvalPlus, stratified LiveCodeBench, +GSM8K, MATH-500, WikiText prompt PPL, tool-call validity, and long-output coding +health. At least three predeclared sampling seeds are required for executable +coding scores. No suite may show a statistically supported regression versus +the same NVFP4 target-only route. DFlash2 request-mean acceptance may not fall +more than `0.05` from the accepted dense DFlash2 reference. + +## Performance gate + +Quality fixes are measured on the same practical 256K/4096/prefix/tool contract. +The accepted short baseline is approximately 18 ms per complete speculative +round; high-acceptance prompts can exceed 220 output token/s. A safe change +must preserve the retained grouped long-context verifier and should keep the +complete-round regression below 1%. Candidate-order QPN8 is not retained merely +to save tens of microseconds if dense-order or dense-selector scoring is more +stable. + +## 2026-08-26 bounded multi-seed isolation result + +The same physical V100 group 4--7 ran all five arms sequentially. The bounded +screen replays two GSM8K, two MATH-500, two HumanEval, and two embedded MBPP +rows at three predeclared request seeds, for 24 outputs per arm. Sampling is +temperature 1.0/top-p 0.95/top-k 20 with `xhigh` reasoning, natural EOS, and a +2,048-token cap. The embedded MBPP rows are useful only for relative trajectory +health: both target-only and every DFlash arm score 0/6 because this cap ends +reasoning before valid code or because the embedded prompt lacks the corrected +independent-test contract. They are not an absolute coding score. + +| Arm | Valid score pattern | Total | Natural stop | Request acceptance | Mean steady decode | +| --- | --- | ---: | ---: | ---: | ---: | +| `T0` | GSM 5/6, MATH 3/6, HumanEval 6/6 | 14/24 | 15/24 | n/a | 75.246 tok/s | +| `D0` | GSM 6/6, MATH 3/6, HumanEval 6/6 | 15/24 | 16/24 | 4.64385 | 237.849 tok/s | +| `D1` | GSM 5/6, MATH 3/6, HumanEval 6/6 | 14/24 | 16/24 | 4.64085 | 244.086 tok/s | +| `D2` | GSM 5/6, MATH 3/6, HumanEval 6/6 | 14/24 | 16/24 | 4.62132 | 253.552 tok/s | +| `D3` | GSM 4/6, MATH 3/6, HumanEval 6/6 | 13/24 | 15/24 | 4.67840 | 259.833 tok/s | + +`T0` and `D1` have exactly the same 24-row pass/fail set despite zero equal +token hashes. This is direct evidence that random trajectory divergence alone +is not a quality regression. `D0 -> D1` changes request-mean acceptance by only +`-0.002999` while improving mean steady decode by 2.62%; compact rejection is +retained. `D1 -> D2` exchanges one win in each direction, keeps aggregate score +unchanged, changes acceptance by `-0.019532`, and improves mean steady decode +by 3.88%; dense-order selector QPN8 remains the quality-first accelerated +candidate. + +`D2 -> D3` yields one D2-only pass and no D3-only pass in this small screen, +reducing 14/24 to 13/24 for a 2.48% mean-decode gain. McNemar `p=1.0` and the +wide bootstrap interval do not prove a population regression, but they also do +not authorize the packaged candidate-order override. Combined with retained +real-hidden evidence, the quality/risk trade is unfavorable. The source now +ignores `VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER=0` unless the separate explicit +benchmark-only `VLLM_SM70_DFLASH2_QPN8_ALLOW_CANDIDATE_ORDER=1` is present. +This converts stale production launch files to the D2 path while preserving +the main QPN8 speedup and leaves D3 available only for controlled research. + +This screen loaded stale stable-ABI RMSNorm SHA `fe986a...9ba8c`; the runtime +capability guard correctly selected the bitwise per-layer fallback. Therefore +its score evidence is valid and all arms are comparable, but its throughput is +directional rather than a practical-endpoint claim. The formal API and speed +gates force repaired SHA `7689e5...d449`. + +## Structured-output safety backport + +The local scheduler previously derived the new reasoning window from stale +async placeholders, fed the whole boundary step to the grammar, advanced FSM +state through drafts produced before their masks existed, and allowed +XGrammar to consume tokens after termination. Those are state-machine bugs, +not model-sampling variation. The targeted backport now: + +- uses the exact `new_token_ids` appended by the scheduler step; +- records the reasoning-end token and advances only the suffix after it; +- grammar-validates post-marker drafts before temporary bitmask-state advance; +- keeps the bonus row constrained after invalid `-1` draft padding; +- stops XGrammar acceptance/validation at termination and resets termination + state explicitly. + +The combined DFlash2, reasoning-structured-output, XGrammar termination, and +quality-comparator gate first passed `110` tests. The final CPU-only rerun, +which deliberately hides GPUs while another process owns them, passes `106` +and skips 12 explicit CUDA cases. The separate V100 compact-rejection gate +passes all eight temperature 0.6/1.0 cases. The two exact failure shapes +reported upstream pass without `Failed to advance FSM`. + +The repaired-extension TP4 Graph API gate then ran with prefix caching, Mamba +align, `qwen3` reasoning, and `qwen3_coder` tools. At B1 and B4 it repeated +`json_object`, strict JSON schema, and required `get_weather` tool calls four +times each: both lanes pass `12/12`, all 24 requests finish successfully, and +the server log contains no FSM-advance, traceback, EngineCore error, HTTP 500, +or `FINISHED_ERROR` signature. The launch intentionally supplied the stale +`DENSE_ORDER=0` value; the runtime logged that it ignored it and enabled +`dense_order=True`, proving the production safety admission works. These short +prompts did not produce a prefix-cache hit. A second gate therefore uses the +actual 256K/current-Flash contract and alternates two long ALPHA/BETA prefixes. +All five requests return the exact per-request sentinel and checksum. Prefix +metrics record 59,396 queried tokens and 32,960 hit tokens, far above the +1,024-token admission threshold. The server log remains free of FSM, HTTP 500, +EngineCore, traceback, and `FINISHED_ERROR` signatures. + +The official Qwen model card also distinguishes general thinking sampling from +precise coding: the latter recommends temperature 0.6 with top-p 0.95/top-k 20 +and no presence penalty. Temperature 1.0 remains the general-thinking and +distribution-equivalence contract. Temperature 0.6 is evaluated below as a +separate client-visible precise-coding profile; sampling profiles must not mask +a verifier or grammar defect. + +## Practical 16K coding gate + +The quality-first D2 route completed MBPP-32, HumanEval-32, and stratified +LiveCodeBench-16 at three predeclared seed bases. Each request uses the +practical 256K/4096/prefix/Mamba-align/Graph engine contract, temperature +1.0/top-p 0.95/top-k 20, `xhigh` reasoning, and a 16K natural-EOS cap. MBPP +excludes the one upstream EvalPlus row without a valid independent-test +contract, leaving 31 scored rows per seed. + +| Suite | Samples | Base/score | Plus | Natural stop | Aggregate output | Mean steady decode | Pooled/request acceptance | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| MBPP | 96 (93 scored) | 89/93 | 80/93 | 95/96 | 213.539 tok/s | 236.902 tok/s | 4.061/4.318 | +| HumanEval | 96 | 94/96 | 92/96 | 91/96 | 208.978 tok/s | 245.645 tok/s | 3.972/4.476 | +| LiveCodeBench | 48 | 33/48 | n/a | 32/48 | 170.283 tok/s | 196.541 tok/s | 3.277/3.694 | + +MBPP Base/Plus by seed are `31/29`, `27/24`, and `31/27` out of 31. +HumanEval Base/Plus are `31/30`, `31/31`, and `32/31` out of 32. The middle +MBPP seed demonstrates substantial sampling variance even though its +acceptance and speed are not worse; this is why a single output or acceptance +mean is not an intelligence metric. + +The first seed exactly matches the historical no-DFlash request-seed contract. +Across MBPP plus HumanEval, both routes score Base `62/63` and Plus `59/63`. +The accelerated route gains one Base and two Plus rows on MBPP, while losing +one Base and two Plus rows on HumanEval. Thus the matched aggregate shows no +net score regression, although the changed sampled trajectory moves individual +failures. MBPP/HumanEval expose six 16K cases among 192 outputs, with three +still extractable and correct. LiveCodeBench adds 16 length cases among 48 +outputs. The temperature-0.6 follow-up below improves executable score but not +natural-stop behavior, so it is retained as an optional precise-coding profile +rather than a forced API default. + +LiveCodeBench scores `11/16` for every seed. Aggregated difficulty scores are +easy `18/18`, medium `12/18`, and hard `3/12`. The historical no-DFlash and +old-DFlash first-seed controls also score `11/16`; the current route exchanges +one passing task with each but adds one natural stop (`11/16` versus `10/16`). +Across all three current seeds, the first two have 11 stops and the third has +10. Every failure is a 16K length case; all stopped cases in the first two +seeds pass. This isolates excessive thinking as the remaining coding-product +risk, rather than verifier corruption or a score regression. + +## Precise-coding sampling and PPL distribution gate + +The same practical D2 engine contract was repeated at temperature 0.6 for the +predeclared middle seed, which was the weakest temperature-1.0 MBPP +realization. All other engine, dataset, prompt, seed, reasoning, and 16K +natural-EOS settings were unchanged. + +| Suite | Temperature 1.0 | Temperature 0.6 | Natural stop, 1.0 -> 0.6 | +| --- | --- | --- | ---: | +| MBPP | Base 27/31, Plus 24/31 | Base 29/31, Plus 27/31 | 31/32 -> 30/32 | +| HumanEval | Base 31/32, Plus 31/32 | Base 31/32, Plus 31/32 | 30/32 -> 30/32 | +| LiveCodeBench | 11/16 | 11/16 | 11/16 -> 10/16 | + +Across the 80 requests, temperature 0.6 changes request-mean acceptance +`4.27345 -> 4.51356`, pooled acceptance `3.70902 -> 3.84470`, mean steady +decode `233.187 -> 244.520 token/s`, and output-token/decode-time throughput +`195.817 -> 201.852 token/s`. It improves MBPP executable scores and lowers no +suite's aggregate score, but natural stops fall from `72/80` to `70/80`. +LiveCodeBench exchanges one pass in each direction: one temperature-1.0 stop +becomes a temperature-0.6 length failure while a different 16K length output +becomes executable and passes. This is sampling variation, not evidence that +temperature 0.6 fixes long-thinking termination. Therefore temperature 0.6 is +accepted as an optional precise-coding profile, not as a forced global API +default; the general-thinking profile remains temperature 1.0. + +A separate current-source Graph PPL pair compares target-only and DFlash2 on +eight fixed 2,048-token WikiText segments. Both arms use the NVFP4 target, +E5M2 target KV, TP4 V100, 256K maximum context, 4096 chunking, prefix caching, +Mamba align, and the repaired extension set. The 16,376 scored prompt tokens +produce weighted target-only/DFlash2 PPL `5.4993116/5.4993622`, an absolute +change of `+0.0000506` (`+0.00092%`). Maximum per-segment PPL difference is +`0.0062143`; prompt-logprob mean/max absolute differences are +`0.0048521/0.481627`. The predeclared `0.01` maximum segment-PPL gate passes +with no failed evidence. The generic comparator labels the artifact +`B-pending` only because this prompt-PPL experiment intentionally does not +collect output-logprob or sampler-logit tensors; the scored generation, +compact-rejection equivalence, and API state gates cover those separate +behaviors. The aggregate PPL result rejects a distribution-level degradation +of the target model from enabling DFlash2. + +## Current status + +The source worktree is isolated. The five-arm screen, safe selector admission, +minimal structured-output backport, comparator, focused CPU/GPU tests, actual +structured API gate, three-seed practical coding gate, 256K long-prefix state +gate, temperature-0.6 repeat, and current Graph PPL pair are complete. The +benchmark supports index-derived request seeds so its first lane exactly +matches the historical target-only seed contract, and it records suite +metadata for official EvalPlus/LiveCodeBench scoring. Compact rejection +matches dense rejection at both temperature `1.0` and precise-coding +temperature `0.6`. The same-card bounded D1/D2 screen already attributes a +3.88% steady-decode gain to quality-safe dense-order selector QPN8; a larger +practical D1/D2 performance pair is optional follow-up rather than a quality +blocker. Development and Draft PR #18 are private; no public branch is used. diff --git a/docs/design/sm70_glm53_flash_nvfp4.md b/docs/design/sm70_glm53_flash_nvfp4.md new file mode 100644 index 0000000000..975b709b1d --- /dev/null +++ b/docs/design/sm70_glm53_flash_nvfp4.md @@ -0,0 +1,176 @@ +# SM70 GLM-5.3-Flash NVFP4 Bring-Up + +Date: 2026-08-27 + +## Contract + +- Checkpoint: `LibertAIDAI/GLM-5.3-Flash-NVFP4`, revision + `9e0d74e3cef17f634e84fb8e2223707e02616290`. +- Hardware: eight V100 GPUs, TP4 and PP2. PP is required for weight capacity; + every PP stage contains one four-GPU TP group. +- Runtime activations and unquantized weights: FP16 on SM70. The checkpoint's + BF16 tensors are cast during load because V100 has no native BF16 execution. +- Speculative decoding: disabled. The checkpoint contains one NEXTN/MTP layer, + but MTP results do not count toward this acceptance. +- Speed gate: batch-one text generation, exact 1024-token input and 256-token + output, steady pure decode reported separately from TTFT and prefill. The + initial target is at least 70 output tokens/s with CUDA graphs enabled. +- Quality gate: fixed-seed greedy token identity across three same-process + requests, finite logits, coherent Chinese and English output, and a natural + official-sampling completion. + +## Architecture + +The language backbone has 45 layers, hidden size 4096, and a 1,048,576-token +position limit. Attention follows an exact four-layer cycle: + +- 34 KDA linear-attention layers at indices `0,1,2,4,5,6,...,44`. KDA uses + 64 heads of dimension 128, three independent convolutional states with a + kernel width of four, and a recurrent FP32 matrix state. The safe gate has a + configured lower bound of -5. +- 11 DeepSeek-style sparse MLA layers at indices `3,7,11,...,43`. They use + NoPE MLA with 64 query heads, 256 NoPE Q/K dimensions, 256 value dimensions, + a 1536 Q LoRA rank, and a 512 KV LoRA rank. +- Every sparse layer has a 32-head, 128-dimension DSA indexer. It selects 2048 + pools, compresses four tokens into each KPool entry, and always retains the + incomplete tail pool. The indexer and tail require separate scheduler + semantics even though the tail co-owns the indexer's backing tensor. + +The residual path uses mHC4 with 20 Sinkhorn iterations. Layers 0-2 have dense +12288-wide SwiGLU MLPs. Layers 3-44 have 288 routed experts, top-8 routing, a +2048-wide expert intermediate, one separate shared expert, sigmoid routing, +FP32 router logits, normalized top-k probabilities, and a routed scale of 2.5. + +The model is natively multimodal. Its vision tower has 24 layers, hidden size +1024, 16 heads, 448-pixel images, 14-pixel patches, temporal patch size two, +and a 4096-dimensional language projection. Initial SM70 acceptance is for the +text path; image and video serving need their own memory and quality gates. + +## Checkpoint Format + +The checkpoint contains 120 safetensor shards and is about 181 GiB. It was +produced by ModelOpt 0.45.0. Its `NVFP4` configuration is weight-only: +`input_activations` and `output_activations` are null. Only the routed-expert +gate/up/down matrices are packed as E2M1 FP4 with FP8 E4M3 block scales at +group size 16 and FP32 tensor scales. Attention, KDA, the sparse indexer, +shared experts, routers, dense MLPs, vision, MTP, mHC, embeddings, LM head, +and norms remain BF16. + +Under TP4, each routed expert uses the validated local shapes: + +- packed gate/up: `[288, 1024, 2048]` bytes after stacking; +- packed down: `[288, 4096, 256]` bytes after stacking; +- global contract: hidden 4096, intermediate 2048, 288 experts, top-8. + +The SM70 path combines block and global scales once, repacks directly into the +TurboMind layout, and deletes the checkpoint tensors. It never keeps a full +FP16 expert-weight copy. + +## Upstream Disposition + +- vLLM PR [#53906](https://github.com/vllm-project/vllm/pull/53906) is the + primary model implementation. It is still open, needs a rebase, and mixes + the model with broad scheduler, connector, multimodal, and runner changes. + This tree ports the model and the required hybrid-cache pieces instead of + merging the entire conflicting commit. +- SGLang PR [#36507](https://github.com/sgl-project/sglang/pull/36507) confirms + the same KDA/DSA/KPool/mHC/NEXTN architecture, but its reported validation is + FP8 on 4x GB300 TP4/EP4 and 8x H100 TP8/EP8. It is useful as a semantic + reference, not as an SM70 kernel or speed baseline. +- SGLang PR [#36513](https://github.com/sgl-project/sglang/pull/36513) reports + final-weight Blackwell/Hopper serving recipes and notes that some measured + BF16 high-throughput runs disabled shared-expert fusion. 1Cat keeps the GLM + shared expert as a separate dense branch, so routed NVFP4 experts are not + fused with an incompatible BF16 shared-expert tensor. + +## 1Cat Integration + +The adaptation is divided into independently testable surfaces: + +1. Register `glm5_next`, its Transformers config and processor, multimodal + wrapper, text backbone, KDA, sparse MLA/indexer, mHC, and NEXTN modules. +2. Carry KDA's bounded gate into the local FLA recurrent/chunked operators and + preserve the existing SM70 recurrent schedule. +3. Add GLM hybrid-cache grouping: 34 KDA states are balanced into four Mamba + groups; 11 MLA and 11 compressed-indexer tensors own the physical slots; + the 11 four-token tail caches co-own their sibling indexer allocations. +4. Virtually split the actual TP4 indexer storage block of 288 entries into + nine 32-entry kernel pages. Exclude the one-block-per-request tail scratch + group from prefix-cache hashing and global block-size selection. +5. Carry mHC post/combination state across the PP boundary. The PP payload is + hidden `[B,4096]`, residual `[B,4,4096]`, post `[B,4,1]` FP32, and + combination `[B,4,4]` FP32. +6. Admit ModelOpt NVFP4 only on exact SM70 and route the validated GLM expert + contract to the graph-safe TurboMind kernel. No Marlin or dequantized-FP16 + fallback counts as accepted performance. +7. Encode KPool E4M3 entries in software on exact SM70. V100 cannot execute + native FP8 conversions, so the writer stores the checkpoint-compatible + E4M3 byte representation through `uint8` pointers. Sparse-indexer query + rotation and scoring stay in FP16 and reuse the existing SM70 HMMA path. +8. Store sparse-MLA latent KV in a GLM-specific packed E4M3FN page. Each token + uses 512 data bytes plus eight UE8M0 power-of-two scales (64 values/group), + for 520 bytes total. Decode keeps this packed representation resident and + gathers/dequantizes only the selected latent rows into a fixed-width FP16 + workspace before two Tensor Core GEMMs. The older direct scalar kernel is + retained only as a reference/test path and is not accepted for B1 decode. + Use the explicit `fp8_e4m3` cache dtype because the historical generic SM70 + `fp8` alias resolves to E5M2 for other model families. +9. Keep all GLM mHC4/H4096 execution on native SM70 kernels. Small-M fused + decode follows the DeepSeek-V4 FP32 staging design, but its final Sinkhorn, + residual mix, and RMSNorm stage is a dedicated single-CTA CUDA kernel for + exact FP16 SM70. Standalone pre/post and large-M prefill use dedicated + Triton kernels that avoid TileLang's SM70 BF16-header compilation failure. + Layer zero also uses the mathematically equivalent broadcast weight, + avoiding a four-stream expanded GEMV. +10. Fuse KDA's two B1 `128 -> 2048` f/g projections into one exact-shape SM70 + CUDA launch. It preserves FP32 accumulation and FP16 stores, is graph-safe, + and hard-errors if the native operator is absent for the accepted TP4 + contract instead of silently returning to two cuBLAS launches. + +## Validation Status + +The 120 checkpoint shards are complete at the pinned revision (about 181 GiB). +The full source extension build completed on the target host. Import checks +confirm the SM70 NVFP4 prepare/dense-stage operators, strided MoE pointer +builder, and MoE permutation operator are registered. Static model/config +parsing, PP intermediate construction, weight-name mapping, KPool E4M3 byte +encoding, KPool tail slot mapping, sparse decode sequence lengths, hybrid +cache layout, and compressed physical-page reshape are covered by focused +tests. The KPool GPU suite passes 28 tests, and the current CPU integration +suite passes 64 tests. + +An eight-process Gloo construction smoke using the exact TP4/PP2 rank layout +also passes on `meta`: PP0 owns 23 layers, PP1 owns 22 layers, every rank +selects `ModelOptNvFp4SM70MoEMethod`, and the PP payload carries `hidden_states`, +`residual`, `post`, and `comb`. + +Exact eight-V100 dummy-weight requests now pass in eager mode for both FP16 and +packed E4M3FN KV. Runtime logs select the GLM sparse backend, ModelOpt NVFP4 +TurboMind MoE, KDA recurrent kernels, KPool/indexer kernels, standalone and +fused SM70 mHC kernels, and direct packed-FP8 sparse MLA. At the same cache +budget, reported capacity increases from 66,355 tokens with FP16 KV to 111,957 +tokens with packed E4M3FN KV (about 1.69x). The focused mHC suite passes 17 +tests on V100, including M=1 decode and M=17 prefill boundaries. + +The real checkpoint now loads and serves on eight V100s with TP4/PP2, packed +E4M3FN KV, no MTP, and full decode CUDA Graph capture. A retained 16-input, +16-output route benchmark improved steady pure decode from 5.35 tok/s on the +initial scalar sparse-attention route to 45.83 tok/s after the packed-KV B1 +gather/dequant plus Tensor Core GEMMs. Both runs emitted the identical token +hash `70007811d61c68bb6ec6b4ac5758744f2d4c6b64b51b6b1352f380091c990902`. +This is a short optimization checkpoint, not the 1024/256 acceptance result. + +An Nsight Systems graph-node trace reports a 23.76 ms mean replay interval and +95.09% graph-node coverage. The sparse attention hotspot is gone; the remaining +rank-average GPU service is dominated by pipeline send/receive wait (12.46 ms, +including the idle peer stage), dense GEMV/GEMM (4.55 ms), routed NVFP4 MoE +(1.44 ms), and mHC (1.40 ms). The native 128-thread mHC final kernel benchmarks +at about 16.9 us with 48 registers, 48 bytes of shared memory, and no spills. +The fused KDA f/g projection benchmarks at 6.28 us versus 23.02 us for two +cuBLAS calls, a 3.66x local speedup and an estimated 0.57 ms per model token. +Its real-model end-to-end rerun remains pending an uncontended eight-GPU slot. + +The exact 1024-input/256-output, three-repeat speed gate and broader Chinese and +English quality gate remain open. Retain all JSON, logs, and profiles under +`/data/minimax-h3/task-cache/glm53-nvfp4-sm70-20260827`; a route-hit smoke or +short quality completion must not be reported as the 70-token/s result. diff --git a/docs/design/sm70_nvfp4_turbomind_operator_optimization.md b/docs/design/sm70_nvfp4_turbomind_operator_optimization.md index 5626b03d2e..77c5d96789 100644 --- a/docs/design/sm70_nvfp4_turbomind_operator_optimization.md +++ b/docs/design/sm70_nvfp4_turbomind_operator_optimization.md @@ -640,11 +640,11 @@ Required run shape unless a note says otherwise: - CUDA device: one V100/SM70 rank - verifier shape: `m=17`, matching `budget=16` plus root - weighted verifier suite: - - `63` MLP gate/up calls - - `63` MLP down calls - - `47` linear-attention qkv calls - - `47` linear-attention z calls - - `16` full-attention o-proj calls + - `63` MLP gate/up calls + - `63` MLP down calls + - `47` linear-attention qkv calls + - `47` linear-attention z calls + - `16` full-attention o-proj calls - time only `awq_gemm_sm70_out` after `awq_sm70_prepare` and warmup - record exact env, JSON artifact, selected kernel when tracing is enabled, and whether the candidate is default behavior or env-forced only diff --git a/docs/design/sm70_qwen38_flash_next_nvfp4.md b/docs/design/sm70_qwen38_flash_next_nvfp4.md new file mode 100644 index 0000000000..3027d31caf --- /dev/null +++ b/docs/design/sm70_qwen38_flash_next_nvfp4.md @@ -0,0 +1,509 @@ +# Qwen3.8 Flash Next NVFP4 on SM70 + +## Status and ownership + +- Status: source bring-up, pinned-host PLE, native Qwen4Exp MTP, and + prefix-cache configuration are implemented; focused CPU/configuration gates + pass and the local ModelScope snapshot is fully verified. The current + accepted no-MTP TP4 candidate reaches 82.274 steady decode tokens/s on the + 8192x512 contract, with exact arithmetic and Chinese quality gates. Its + direct NVFP4 ten-route QPN-M1 MoE route is now the default for the exact + SM70/TP4/B1 contract. The preceding 77.370-token/s graph-node trace remains + the detailed optimization control. Native MTP4 also completes TP4 model + load, graph + capture, warmup, and two 1024x256 requests; its separate acceptance and + cycle-cost evidence remains recorded below. +- Integration line: public `origin/main`, through Draft PR + [#345](https://github.com/1CatAI/1Cat-vLLM/pull/345). +- Base SHA: `d63e9490f65f9e01f6649053c1ab72922034b931`. +- Model: `RadixArk/Qwen3.8-Flash-Next-NVFP4` at revision + `7b719225242aacd3dbd3f9407468c2ee9a9d2594`. +- Model download: `/data/models/RadixArk/Qwen3.8-Flash-Next-NVFP4`. +- Download source: ModelScope `master`, verified against the fixed Hugging Face + revision above: all 419 file sizes match and all 208 comparable LFS SHA-256 + values match. After download, all 419 local files were independently hashed + against the ModelScope manifest with zero missing, size-mismatched, or + SHA-256-mismatched files. +- Upstream references: + [vLLM PR 53896](https://github.com/vllm-project/vllm/pull/53896) and + [SGLang PR 36497](https://github.com/sgl-project/sglang/pull/36497). + +The upstream PRs are implementation references, not acceptance evidence. Both +were still open when this work started, so imports must be narrowed to the +Qwen4Exp route and validated against this branch. + +## Frozen first-pass contract + +The first correctness route deliberately excludes speculative decoding. + +- Hardware: four NVIDIA V100-SXM2-32GB GPUs (SM70). +- Parallelism: TP4, PP1, no expert parallelism. +- Model mode: `--language-model-only`. The initial SM70 route deliberately + excludes the vision tower from its memory, quality, and performance gates; + omitting the flag fails during configuration instead of reaching a private + Qwen3.5 multimodal API mismatch at model construction. +- Compute dtype: FP16; no BF16 or native FP8/NVFP4 tensor-core assumptions. +- Checkpoint: ModelOpt NVFP4 routed-expert weights, consumed as an SM70 + weight-only W4A16 route. Ignored dense, attention, GDN, shared-expert, GR, + PLE, and LM-head modules stay in their checkpoint dtypes and execute in + FP16 where required. +- PLE/N-gram table: allocate and load each TP shard directly in pinned host + memory. It must never be materialized on a GPU before being moved to the + host. Gathered rows are transferred asynchronously and converted to FP16. +- Initial KV cache: FP16. FP8 KV cache is a separate, quality-gated follow-up. +- Initial decoding: MTP disabled. MTP may be enabled only after the no-MTP + route is correct and its emitted-token baseline is recorded. +- Model runner: V2 is the default initial route. Its Qwen4Exp model state keeps + raw token IDs and builds the PLE context from committed tokens, so rejected + speculative candidates cannot leak into the next trigram. V1 remains a + correctness control, not the primary performance route. +- Prefix caching: enabled for the MTP validation route. Hybrid recurrent state + uses `mamba_cache_mode=align` with chunked prefill. The fixed QSA compressor + ring is explicitly non-cacheable and is excluded from prefix-hit + reconciliation; a clean ring block is allocated after a hit. Main QSA KV, + compressed QSA KV, and aligned GDN/Mamba state remain cacheable. + +The native-MTP validation route keeps the same TP4/PP1, FP16 activation, +ModelOpt NVFP4, language-model-only, and FlashAttention-V100 contract. It uses +four speculative tokens, V2, CUDA graphs, prefix caching, a deterministic +greedy prompt, and two identical requests. The second request is both the hot +speed sample and the prefix-cache reuse check. A matched no-MTP run is still +required for exact token-ID quality comparison and incremental verifier cost. + +## Architecture facts that affect the port + +The text stack has 48 layers: 36 gated-delta-net layers and 12 QSA layers in a +3:1 pattern. Hidden size is 2560. QSA uses 24 query heads, two KV heads, +head-dimension 256, index dimension 128, compression ratio four, and a 2048 +token sparse budget. The MoE has 512 routed experts, top-10 routing, a 640-wide +routed expert, and one 640-wide shared expert. General residual connections +use four streams and rank 320. + +PLE is a learned trigram embedding, not prompt-ngram speculative decoding. It +uses 16 heads (`ngram_size=3`, eight heads per n-gram order), embedding width +2560, and FP8 E4M3 storage. Native speculative decoding is the separate +one-layer MTP head. + +## Memory budget hypothesis + +Safetensor payloads total about 125.910 GiB. The sharded PLE payload is about +47.684 GiB, or about 11.921 GiB of pinned host memory per TP rank. Removing PLE +from device residency leaves an idealized 19.556 GiB of checkpoint payload per +GPU before replicated tensors, KV/index caches, CUDA graphs, and workspaces. + +This is a planning bound, not a measured peak. Startup must record host RSS, +pinned memory, per-rank device peak, post-load device residency, and whether a +loader creates duplicate staging buffers. A 262144-token context is admitted +only after the measured peak leaves a safe margin on every 32GB GPU. + +The SM70 TurboMind repack changes routed-expert FP4 scales from FP8 to FP16. +For TP4, routed experts are estimated at about 15.82 GiB/rank in the source +checkpoint and 17.57 GiB/rank after repack. This puts the idealized final +device weights near 21.3 GiB/rank. Because layers are repacked sequentially, +the estimated transient weight peak is about 22.1 GiB/rank before runtime +buffers, KV/index caches, NCCL, and CUDA graphs. These are storage calculations, +not `torch.cuda.max_memory_allocated` measurements. + +The loader marks the PLE parameter as permanently host-resident so generic +quantization post-processing cannot stage the entire 11.921 GiB TP shard on a +GPU. Only its small scale parameter resides on device; lookup reads selected +FP8 rows through a stable UVA view and converts the gathered output to FP16. + +With the real SM70 platform alignment and the QSA attention backend selected, +the V2 scheduler block is 784 tokens, the recurrent-state block is 32768 +tokens at a 32768-token initial maximum length, and each padded recurrent page +is 802816 bytes. The exact synthetic model layout has one 24-layer uniform QSA +main/compressed group, one 12-layer fixed circular QSA ring group, three +12-layer GDN state groups, and one PLE short-convolution state group. It +allocates 24 physical cache tensors. The aligned pool cost is 10235904 bytes +(9.762 MiB) per shared block per TP rank. + +For one request, the resulting cache-pool planning estimates are about 0.448 +GiB/rank at 32K (47 shared blocks), 1.649 GiB/rank at 128K (173 blocks), and +3.241 GiB/rank at 262144 tokens (340 blocks). Combining the last figure with +the estimated 21.3 GiB final weights gives about 24.54 GiB/rank before CUDA +graphs, workspaces, NCCL, allocator fragmentation, and loader transients. This +explains why TP4 is plausible, but it is not evidence that the maximum context +will load safely. + +## No-MTP TP4 decode control and architecture audit + +The current unprofiled control is +`.artifacts/qwen38_flash_next_nvfp4_20260826/results/target80-qpn-m1-mtp-off-8192x512.json`. +It uses one request, TP4/PP1, FP16 activations, MTP off, V2, CUDA graphs, the +direct QPN-M1 NVFP4 experts, online channel-QPN8 attention/GR projections, +FlashAttention-V100 plus FlashQLA, and pinned-host PLE. Three matched repeats +measure 82.2521, 82.2750, and 82.2956 steady decode tokens/s: mean 82.274268 +tokens/s, or 12.154469 ms/token. The repeat range is only 0.043466 tokens/s. +This is 4.904026 tokens/s faster and 0.770399 ms/token lower than the preceding +77.370242-token/s architecture control. The arithmetic gate emits `42` with +the unchanged token hash `93def17b...`, and the Chinese gate emits the expected +statement that water boils at 100 degrees Celsius under standard atmospheric +pressure with the unchanged token hash `11d98c01...`. Peak sampled device +memory is 29,752 MiB on each 32,768-MiB V100. Model load reports 20.86 GiB/rank, +3.73 GiB/rank available for KV cache, and a 0.16-GiB/rank graph footprint. + +The graph-node trace is +`.artifacts/qwen38_flash_next_nvfp4_20260826/profiles/qwen38_nvfp4_target80_arch_mtp_off_i8192_o16_graph_nodes.nsys-rep`; +its stable parsed report is the adjacent `*_per_token_steady.{md,csv,json}` +set. Mean wall TPOT is 14.046 ms under Nsight, rank-max GPU service is 13.797 +ms, and the residual host/scheduler gap is only 0.248 ms. Each rank replays +about 1,928 graph kernels per token, so both memory traffic and launch count +matter. + +| architecture path | rank-average service | important sub-costs | +|---|---:|---| +| GR/HC read, mix, and combine | about 2.46 ms | QPN8 down 1.047 ms, reduce/SiLU 0.250 ms, QPN8 up+gate+mean 0.763 ms, combine+norm 0.400 ms | +| QSA selection and sparse attention | about 1.08 ms | MQA scorer 0.289 ms, persistent top-k 0.214 ms, sparse attention 0.444 ms, split merge 0.110 ms | +| attention input/output projections | about 1.27 ms | split-16 QPN8 0.789 ms, split-12 QPN8 0.480 ms | +| routed MoE expert GEMMs | 1.675 ms | TurboMind W4A16 expert kernels | +| MoE route and reduction | about 1.07 ms | specialized top-k 0.568 ms, activation 0.287 ms, prepare 0.107 ms, and weighted reduce 0.105 ms | +| TP communication | about 0.87 ms | exact 5-KiB four-rank push all-reduce dominates | +| GDN recurrent core and convolution | about 0.41 ms | recurrent update 0.282 ms and convolution update 0.126 ms | +| LM head and sampling | about 0.70 ms | LM-head GEMV dominates | +| PLE host lookup | about 0.02 ms | already hidden behind the first transformer work | + +This ranking determines the implementation order. It also prevents repeating +architecture work whose maximum possible gain cannot close the current +0.424868-ms/token gap to the first 80 tokens/s gate. + +The accepted QPN-M1 route removes the per-layer routed-input replication and +consumes the existing `(512, 2560, 40)` W13 and `(512, 160, 320)` W2 +TurboMind-packed NVFP4 tensors directly. Across ten checkpoint experts, W13 +falls from 20.983 to 12.213 microseconds with split-K 8; its maximum stage +error is 3.052e-5. W2 split-K 1 is bitwise equal and falls from 8.239 to 7.541 +microseconds. A second screen spans expert IDs 0 through 470 and 16 random +inputs; the final weighted MoE output differs by at most 4.768e-7. Together +with removal of the traced 0.107-ms input-prepare work, the trace projection is +about 12.36 ms/token, or 80.9 tokens/s. The fixed resident-model measurement +above exceeds that projection at 82.274 tokens/s and is the acceptance +evidence; no MTP or speculative decoding is enabled in this result. + +### Upstream methods and SM70 decisions + +The [official Qwen repository](https://github.com/QwenLM/Qwen3.8-Flash-Next) +and its +[technical report](https://github.com/QwenLM/Qwen3.8-Flash-Next/blob/main/tech_report.pdf) +specify the 3:1 GDN/QSA pattern, four-branch GR, and the layer-2 PLE placement. +The report's relevant inference recommendations are fused GR reads/writes, +folding RMSNorm into the read, low-precision residual-branch storage, and +FlashQLA for GDN. It explicitly reports quality loss when GR is sparsified to +only two branches, so branch pruning is excluded from this port. FP8 branch +storage is not copied blindly either: V100 has no native FP8 arithmetic, so a +conversion-inclusive screen is required before it can replace FP16 residuals. + +The official [vLLM Qwen4Exp PR](https://github.com/vllm-project/vllm/pull/53896) +is the structural base for the QSA/GDN/GR model path. The current branch adds +SM70-only projection, cache, attention, and graph routes rather than importing +unrelated upstream platform changes. The +[SGLang day-zero optimization report](https://staging.lmsys.org/blog/2026-08-26-qwen-flash-next) +describes the production HC strategy more precisely: low-M Mix uses split-K +GEMMs with SiLU in the down epilogue and sigmoid/four-stream reduction in the +up epilogue, while Combine splits the hidden dimension to expose more CTAs. +The accepted QPN8 HC route already mirrors those two epilogue fusions and the +SM70 Combine kernel already partitions the hidden dimension; importing the +SM100-only CuTeDSL implementation from +[FlashInfer PR 4266](https://github.com/flashinfer-ai/flashinfer/pull/4266) +would therefore add no V100 kernel. The exact V100 screen also rejects +collapsing the accepted two-QPN8-kernel path into one persistent launch. The +[SGLang day-zero HC kernel](https://github.com/sgl-project/sglang/blob/02d38b77db92699e5d4f1a78226bf711e9cc762a/python/sglang/srt/layers/hc_mix_triton.py) +replaces the FP16 down-GEMV, SiLU, up-GEMV, sigmoid, and four-branch mean with +one persistent kernel. A V100 form preserves the existing QPN8 weights and +exact split-K reduction order. All 32 checkpoint-weight seeds and 1,000 graph +replays are bitwise clean, but its cold median is 26.624 microseconds versus +23.552 microseconds for the accepted path. At 97 HC calls per token that would +regress TPOT by about 0.298 ms, so the experiment is removed. Using SGLang's +FP16 weights unchanged would also double the GR weight traffic. +SGLang's fused QSA norm/RoPE/cache-store kernels map to only about 0.063 +ms/token in this trace and are therefore a later cleanup, not the 80-token/s +critical path. + +[TokenSpeed's Qwen4Exp implementation](https://github.com/lightseekorg/tokenspeed/tree/4cb771487f2a070bc3a39c36b6dc8d959a36a92f) +adds a streaming QSA scorer/top-k that avoids materializing the full FP32 score +matrix, as well as FP8 PLE storage. The streaming design is attractive at long +context, but the exact 8K graph has only 2,048 visible compressed blocks out of +a 32,768-block capacity. TokenSpeed's default geometry would launch 64 split +programs and merge 64 by 512 packed candidates while only four splits contain +work. It must beat the already-screened two-warp scorer plus existing top-k in +an exact cold-cache microbenchmark before being ported. This checkpoint's PLE +is already FP8 E4M3 on pinned host memory, so TokenSpeed's FP8 table route does +not reduce its current 11.92-GiB/rank footprint or materially improve the +measured 0.02-ms PLE latency. + +The original [Gated DeltaNet paper](https://arxiv.org/abs/2412.06464) and the +[Flash Linear Attention project](https://github.com/fla-org/flash-linear-attention) +support recurrent decode plus chunked/parallel prefill as the intended +hardware decomposition. This branch already uses FlashQLA for the recurrent +decode state update and keeps FLA-style packed state handling. The current +trace assigns only about 0.41 ms/token to the GDN recurrent core plus short +convolution, versus more than 2 ms to its surrounding GR and projections, so +replacing the GDN recurrence is lower priority than removing those projection +and launch costs. The [Hyper-Connections paper](https://arxiv.org/abs/2409.19606) +also makes clear that the four streams carry distinct learned paths; pruning or +averaging them before their learned gates would change the architecture, not +merely optimize it. + +The currently admitted candidates are deliberately narrow: + +- QSA MQA scorer: 32 columns and two warps for SM70 M=1. The checkpoint-shape + cold-cache screen falls from 29.696 to 11.264 microseconds, without changing + selected block IDs. The projected 12-layer saving is about 0.22-0.33 + ms/token. +- MoE router: an exact M=1, E=512, top-10, FP16, renormalized SM70 kernel. It + falls from 11.729 to 5.657 microseconds in the source screen and preserves + NaN, positive-infinity, all-negative-infinity, and tie behavior. The trace + projection is about 0.53 ms/token. Model-level acceptance also requires the + `SM70 Qwen3.8 E512/K10 router top-k path enabled` marker. +- GR/HC persistent launch: rejected despite bitwise output and clean replay + state because it is 3.072 microseconds slower per call in the checkpoint- + shape cold-cache screen, a projected 0.298-ms/token regression. The accepted + split-QPN8 down/reduce/up implementation remains unchanged. +- MRv2 greedy sampling: the active V2 runner always materialized globally + gathered logits and then launched Gumbel sampling, even for the exact + temperature-zero single-request decode used here. The new SM70-only route + keeps the same full local LM-head calculation, performs the existing + TP-local exact argmax/pair gather, and falls back whenever logprobs, grammar, + penalties, bad words, logit bias, NaN counting, random sampling, prefill, or + speculative decoding is present. This removes the traced 0.103-ms Gumbel + kernel and 0.044-ms full-logit all-gather before accounting for the local + argmax cost. + +The QSA scorer, router, and V2 greedy projections together expose about +0.90-1.01 ms/token of gross trace cost, which is enough on paper to cross the +immediate 80 tokens/s gate after allowing for the pair-gather argmax overhead. +A single resident-model run will validate them together; the model will not be +restarted separately for each micro-optimization. The rejected persistent GR +candidate is not included in that run. + +## Native MTP4 TP4 validation snapshot + +The first complete native-MTP run uses V2, TP4, FP16 activations, ModelOpt +NVFP4 weights, `FLASH_ATTN_V100` for target and draft attention, four draft +tokens, `mamba_cache_mode=align`, prefix caching, chunked prefill, and +FULL+PIECEWISE CUDA graphs. It runs two identical deterministic 1024-token +prompts with 256 forced output tokens each. The artifact is +`.artifacts/qwen4_exp_mtp_tp4_20260827/mtp4_prefix_graph_i1024_o256_r2_hetero_v2.json`. + +- Source HEAD is `d9a39ea434` on the Qwen3.8 worktree branch. The measurement + also sees the worktree's separately owned, uncommitted SM70 kernel changes; + it is bring-up evidence and must be repeated from a clean, pinned source + before becoming release-baseline evidence. +- Target plus MTP weights use 23.16 GiB/rank. The aligned MTP attention block + is 816 tokens with 1.62% recurrent-page padding. Available KV-cache memory + is 4.77 GiB/rank, or 219,942 tokens. Observed peak device memory is + 32,330 MiB on every 32,768-MiB V100; the run therefore has only 438 MiB of + peak device headroom at `gpu_memory_utilization=0.90`. +- PLE remains host-resident at 11.92 GiB/rank. The sampled minimum host + `MemAvailable` is 48.107 GiB and minimum free swap is 247.994 GiB. +- Both repeats emit 256 tokens and are exactly equal token-for-token. The first + request has 9.750-second TTFT and 53.125 steady decode tokens/s. The repeated + prefix has 2.667-second TTFT and 52.187 steady decode tokens/s. Mean steady + decode is 52.656 tokens/s; the first end-to-end request includes prefill and + JIT and is not a decode baseline. +- Across 256 speculative steps, the MTP head proposes 1,024 draft tokens and + 254 are accepted. Mean acceptance length including the target bonus is + 1.9921875. Draft acceptance is 24.8047%; per-position acceptance is + 54.6875%, 26.5625%, 13.28125%, and 4.6875%. +- A strictly matched no-MTP run is required before claiming target-token + equality or a verifier-cost ratio. Older 1024x256 artifacts use the distinct + Qwen3.8-27B checkpoint and are not valid controls for Flash Next. + +### MTP4 cycle-cost split + +The graph-node trace is +`.artifacts/qwen4_exp_mtp_tp4_20260827/mtp4_cost_nsys_i1024_o32_hot_u085_v2.nsys-rep`. +It changes only `gpu_memory_utilization` from 0.90 to 0.85 to leave room for +Nsight Systems; the resulting KV-pool size is not a new throughput baseline. +The uninstrumented 0.90 run above remains the absolute-speed evidence. + +Nsight's PIECEWISE ranges are subgraphs rather than emitted-token boundaries, +so the generic per-token parser must not be used for this trace. Stable cycle +boundaries are instead identified by the first W13 kernel of the unquantized +MTP MoE. The first target TurboMind NVFP4 GEMM after that proposal marks the +start of the next target verification. Fourteen complete steady cycles have +the following wall-time split: + +| phase | mean | median | p90 | share of cycle | +|---|---:|---:|---:|---:| +| MTP4 proposal plus scheduler handoff | 13.477 ms | 13.426 ms | 13.699 ms | 31.2% | +| five-token target verify, reject, and state update | 29.714 ms | 29.710 ms | 29.750 ms | 68.8% | +| complete speculative cycle | 43.191 ms | 43.147 ms | 43.416 ms | 100% | + +At the approximately 2.00-token natural acceptance length, this corresponds to +about 14.86 ms of target verification and 6.74 ms of proposal/handoff per +emitted token, or 21.60 ms total (46.3 tokens/s from the traced cycle). The +five-position verifier is only about 40% utilized: two emitted positions per +five computed positions. These normalized costs, rather than raw cycle time +alone, are the comparison gates for MTP3. + +The proposal phase's rank-average GPU-service sums are led by FP16 GEMM/GEMV +(2.791 ms), other MTP/QSA kernels (2.345 ms), and TP communication (1.337 ms). +The verification phase is led by FP16 GEMM/GEMV (10.357 ms), other target +kernels (5.511 ms), target TurboMind NVFP4 GEMM (3.665 ms), SiLU/gating +(1.910 ms), and TP communication (1.762 ms). GPU-service sums can overlap and +are composition evidence, not additive wall time. + +The current route remains fixed at native MTP4. Before changing verifier width, +the acceptance defect must be corrected and remeasured: the first-position +acceptance near 50% shows that the dominant issue occurs before later draft +iterations. Sampling and bookkeeping micro-optimizations come only after a +matched MTP4 acceptance measurement. The 32-token trace prompt still forces +generation after an early EOS and is therefore cost evidence only, not +representative acceptance evidence. + +### Natural-prompt baseline and MTP4 acceptance root cause + +The completed portion of +`.artifacts/qwen4_exp_mtp_tp4_20260827/mtp4_acceptance_natural6_batch_o256_v4_sm70cpp.log` +is the current natural-prompt MTP4 baseline. Four requests complete before the +fifth prefill triggers an asynchronous illegal-memory access. Immediately +before that failure, mean acceptance length is 2.00, 248 of 992 drafted tokens +are accepted, per-position acceptance is 49.6%, 25.4%, 15.3%, and 9.7%, and +aggregate output throughput is 47.54 tokens/s. These are preliminary +four-request statistics, not a stable six-request result. + +The failure is deterministic at the fifth request in both the original +TileLang GDN prefill route and a separately checked SM70 C++ GDN prefill +route. The failing request schedules 88 new tokens and receives block IDs 89 +and 90 for the first two cache groups after four prior requests have finished. +An isolated 88-token C++ GDN prefill matches the FLA reference and survives 64 +repetitions. Because a known fifth-request failure would discard the complete +benchmark JSON, the first optimization run deliberately keeps the exact four +prompts that already complete in the MTP4 baseline. The 88-token shape versus +prefix-cache/block-lifetime distinction remains a separate minimum-size +stability experiment rather than being mixed into the acceptance comparison. + +Commit `70b63a1e45` adds two independently measurable MTP cost optimizations: + +- The V2 Eagle path now honors `use_local_argmax_reduction`. Greedy MTP no + longer gathers a 248,320-column logits tensor on every draft step; it uses + vocab-shard-local top-1 values and IDs followed by a compact TP reduction. + SM70 MTP defaults to greedy drafting so the fast path is reachable, while + DFlash retains probabilistic drafting. +- Qwen4Exp MTP can reuse the step-0 QSA sparse indices on later draft steps. + The prefill output is compacted to one row per request and steps 1+ skip the + repeated QSA top-k. This follows the upstream vLLM lifecycle hooks and the + SGLang Qwen4Exp default, while exposing + `index_share_for_mtp_iteration=false` as an A/B control. + +The acceptance audit found a separate correctness root cause. This checkpoint +stores the BF16 draft routed experts as two fused tensors: +`mtp.layers.0.mlp.experts.gate_up_proj` with shape `(512, 1280, 2560)` and +`mtp.layers.0.mlp.experts.down_proj` with shape `(512, 2560, 640)`. The private +tree's legacy `FusedMoE` recursive mapping recognized only per-expert names +such as `experts.0.gate_proj.weight`. Its child loader consumed the unmatched +generator without reporting those tensors, leaving `w13_weight` and +`w2_weight` at their `torch.empty` initialization. That is about 1.17 GiB of +uninitialized draft parameters per TP rank, or 4.69 GiB across TP4, and directly +explains why first-position acceptance was already abnormally low. + +Commit `cf5c44a1aa` ports the upstream fused-checkpoint aliases for Qwen4Exp and +tests gate/up splitting plus complete expert fan-out. Commit `ceb543c055` makes +the draft loader fail closed if either routed-expert parameter is absent, so a +future mapping regression cannot masquerade as a successful model start. The +focused loader gate is 21 passed tests plus Ruff and format checks. + +The next GPU gate is one TP4, prefix-cache-on, four-prompt native MTP4 run with +greedy local argmax and QSA index sharing explicitly enabled. It uses the same +prompts, order, sampling, and output cap as the completed four-request baseline +and records per-request and per-position accepted length, code output quality, +emitted tokens/s, and peak memory. MTP3 is not an optimization candidate for +this gate. A subsequent Nsight trace is warranted only after the corrected +MTP4 acceptance result is stable; it will report proposal time, five-token +target verification time, and cycle cost per emitted token separately. + +## Acceptance gates + +1. Static route: Transformers config, model registry/processor registration, + QSA/GDN/GR/PLE modules, and ModelOpt NVFP4 mapping load without importing an + Ampere-only backend. Multimodal execution is outside this first route. +2. Loader route: TP4 expert shards select TurboMind SM70 W4A16; PLE shards are + born on pinned CPU memory and do not consume persistent device memory. +3. Numerical route: focused operator comparisons against FP32/FP16 references, + followed by deterministic token-ID and output checks on the full model. +4. Memory route: 32K bring-up first, then 128K and the exact 262144 boundary; + report controlled OOM separately from corrupted output. +5. Performance route: one request, TP4, PP1, MTP off, FP16 activations, 8192 + input tokens and 512 output tokens. Report TTFT/prefill separately and + calculate steady pure decode from emitted tokens 33-512. The first recovery + gate is at least 80 emitted tokens/s (at most 12.5 ms/token); the final target + remains at least 100 emitted tokens/s (at most 10 ms/token), both with CUDA + graphs enabled. Record an otherwise identical eager control. +6. MTP follow-up: report accepted length, target passes, emitted tokens/s, and + output quality separately; do not compare accepted candidates with emitted + tokens. + +## Initial implementation boundary + +Reuse the upstream Qwen4Exp Python structure and tests where they match this +tree. Do not import unrelated AMD, SM90, build-system, or broad engine changes. +The first SM70-specific changes are limited to genericizing the existing +TurboMind NVFP4 MoE shape contract, adding the QSA/indexer route, and adding a +pinned-host PLE loader/gather path. Optimize GDN, GR, sparse attention, and MTP +only after profiles identify them as measured decode bottlenecks. + +## Source validation snapshot + +- The ModelScope download completed successfully at the path above. A full + post-download verification checked 419 files and about 125.910 GiB of + safetensor payload: zero files were missing and zero size or SHA-256 values + differed from the remote manifest. +- The real downloaded `config.json` resolves without remote model code as + `Qwen4ExpConfig` / `Qwen4ExpTextConfig`: 48 layers, 36 GDN, 12 QSA, 512 + experts, top-10, HC count four/rank 320, and one trigram PLE layer. +- Exact-SM70 configuration construction with FP16, TP4, prefix caching off, + language-model-only mode, and V2 selects `ModelOptNvFp4Config`, the + pinned-host PLE default, and the Qwen4Exp PLE/QSA compilation split + operators. The same real configuration rejects the unvalidated multimodal + route with an actionable `--language-model-only` error. +- Exact real-checkpoint construction with native MTP4 and prefix caching on + resolves the target as `Qwen4ExpForConditionalGeneration`, the draft as + `Qwen4ExpMTP`, target and draft attention as `FLASH_ATTN_V100`, and recurrent + caching as `align` with 16-token configured blocks and chunked prefill. + Focused prefix-cache/QSA coordinator and model-config tests pass. The + non-cacheable QSA compressor ring is skipped during prefix-hit matching, + matching the current upstream Qwen4Exp contract. +- TP4 loaded all 206 checkpoint shards in 109.27 seconds and then correctly + rejected an incomplete runtime extension set: `_C` and + `_C_stable_libtorch` were present, but `_moe_C` was omitted. A complete + runtime must carry all three. The matching `_moe_C` artifact has SHA-256 + `a14eeb4fa06947e335cf69ee188e23509fc294da61cada786df09888ca5b4469`; + its graph-safe permute, workspace-size, unpermute schemas, and SM70 support + probe all pass before the next full-model attempt. +- Full 48-layer meta construction from the real checkpoint config succeeds in + language-model-only mode. It instantiates QSA, GDN, HC, PLE, and all 512 + experts without materializing weights; the routed experts select + `ModelOptNvFp4SM70MoEMethod(use_a16=True)` and the PLE table has shape + `(320001536, 160)` with FP8 E4M3 storage. This constructor probe used TP1; + TP4 selection and expert geometry are covered separately. Full TP4 target + plus native-MTP loading now completes with 23.16 GiB/rank of loaded model + state before cache and graph allocation. +- Focused CPU tests cover PLE shard loading and hashing, `seed=None`, permanent + host residency during post-load processing, QSA cache grouping, V1 and V2 + n-gram inputs, V2 circular block-table sizing, scheduler-manager conversion, + official checkpoint weight mappings, and Qwen3.6/Qwen3.8 NVFP4 route + selection. The current CPU-only focused run is 76 passed and 7 CUDA skips; + all 55 changed Python files pass Ruff, format, and compileall checks. +- In the pre-final real V100-SXM2-32GB snapshot, 63 focused tests pass. They + include the Triton V2 slot-mapping kernel with its QSA circular group + disabled, pinned-host FP8 lookup through a CUDA UVA view, the compressed QSA + storage-page reshape, V2 committed-token PLE state, and the SM70 ModelOpt + NVFP4 selection gates. A final V100 rerun is still required after the current + GPU owners release a device. +- The upstream QSA fused pre-indexer executes on SM70 for both ordinary RoPE + and MRoPE inputs and matches a PyTorch normalization reference. This also + exposed and fixed two private-tree API differences: QKV projection is local + because the branch's `Qwen3NextAttention` has no `_project_qkv_gate`, and its + `triton_mrope` accepts eight rather than nine arguments. +- Actual SM70 platform alignment produces a 784-token attention block and an + 802816-byte padded recurrent page; the exact synthetic 48-layer cache layout + validates successfully after that alignment. +- The HC grouped norm/gate/combine kernels pass FP16 reference checks on a real + V100. A captured pinned-host FP8 embedding probe (228.9 MiB synthetic table, + 16 rows by 160 elements per replay) measured 95.81 microseconds/replay, + including the input-ID copy. This only demonstrates that the isolated UVA + lookup can be captured and is not by itself an end-to-end throughput result + or a measurement of the full 11.921 GiB TP shard. +- The existing general KV-cache utility/manager suites pass 69 tests; one + unrelated DeepSeek-v4 fixture failure is unchanged from the integration + base because its `SimpleNamespace` omits `max_in_flight_tokens`. diff --git a/docs/design/sm70_qwen38_fp8_prefill_5500.md b/docs/design/sm70_qwen38_fp8_prefill_5500.md new file mode 100644 index 0000000000..8cf39d62ad --- /dev/null +++ b/docs/design/sm70_qwen38_fp8_prefill_5500.md @@ -0,0 +1,142 @@ +# Qwen3.8 FP8 exact-8K prefill 5500 plan on V100 + +## Scope and acceptance contract + +This work targets one narrow, reproducible contract: + +- model: `/home/ymzx/models/Qwen3.8-27B-FP8`; +- hardware: tensor parallel four on four fully connected V100-SXM2-32GB GPUs; +- prompt/output: exactly 8000 prompt tokens and one sampled token; +- dtype/quantization: FP16 activations, FP8 checkpoint, FP8-E5M2 KV cache; +- attention: Flash-V100 full attention and FlashQLA linear attention; +- scheduler: chunked prefill and prefix caching enabled, one sequence, 8096 + maximum batched tokens; +- graph: `VLLM_COMPILE`, `FULL_AND_PIECEWISE`, compile size 8000, capture + sizes 1/2/8000; +- sampling: temperature 1.0, top-p 0.95, top-k 20, seed 20260823; +- communication: TP custom all-reduce enabled and the accepted exact-8K + reduce-scatter/Gemma RMSNorm/all-gather fusion active. + +The acceptance metric is request-wall prompt throughput, not an isolated +kernel rate. The goal is at least **5500 prompt tok/s**, which requires the +steady request wall time to be at most `8000 / 5500 = 1.454545 s`. Pure +prefill time is reported separately. Correctness requires seven official +sampling repeats, unchanged route admission outside this exact contract, and +an explicit rollback switch. + +## Accepted baseline and remaining budget + +The source-matching accepted baseline produced six steady request-wall samples: + +`1.559085 / 1.561460 / 1.560481 / 1.562664 / 1.564697 / 1.564643 s`. + +The median is `1.562062 s`, or **5121.44 prompt tok/s**. The pure prefill +median is `1.547100 s`, or **5170.96 tok/s**. Seven of seven official samples +returned token `[1061]`, text `" This"`. + +Reaching 5500 therefore requires removing about **107.5 ms** from the steady +request critical path, or 6.88% of the current request wall time. With the +current roughly 15 ms request-minus-prefill overhead, the corresponding pure +prefill target is about `1.4396 s`. + +Baseline artifact: + +`/data/minimax-h3/task-cache/qwen38-fp8-prefill-utilization-20260823/results/full_model_cutlass_all_projections_final_source_tp4_fused_graph_i8000.raw.json` + +## Trace localization + +The retained Nsight Systems trace is: + +`/data/minimax-h3/task-cache/qwen38-fp8-prefill-utilization-20260823/profiles/full_model_cutlass_final_source_nsys_node_i8000_gpus4567.sqlite` + +Each rank executed 1757 kernels on one CUDA stream. After the initial request +setup gap, the model region is nearly continuously occupied, so launch-gap +tuning cannot provide the required 107.5 ms. The following values are mean +per-rank kernel service; the range shows rank skew. Kernel service is not +summed across TP ranks. + +| Category | Mean service | Share | Rank range | +| --- | ---: | ---: | ---: | +| gate/up exact CUTLASS | 468.787 ms | 30.27% | 457.594-476.271 ms | +| QKV/QKVZ exact CUTLASS | 325.594 ms | 21.02% | 318.374-330.700 ms | +| fused collective + Gemma norm | 274.364 ms | 17.72% | 254.165-301.994 ms | +| down exact CUTLASS | 218.743 ms | 14.12% | 214.052-221.915 ms | +| FlashQLA GDN attention core | 78.876 ms | 5.09% | 77.606-80.162 ms | +| full attention | 76.801 ms | 4.96% | 75.609-77.921 ms | +| FP8 weight dequantization | 31.126 ms | 2.01% | 30.264-32.157 ms | +| interleaved SiLU/mul | 20.101 ms | 1.30% | 19.421-20.800 ms | +| other kernels | 54.302 ms | 3.51% | 53.293-55.450 ms | + +This rules out treating attention or QKV as the only target. Gate/up, QKV, +down, and the collective boundary account for about 83% of per-rank GPU +service. The first architecture experiment therefore joins communication and +the following projection instead of tuning either in isolation. + +## Architecture candidate: token-shard consumer projection + +The accepted fused boundary currently reduce-scatters token rows, performs the +mixed-dtype Gemma RMSNorm on the owned 2000-row shard, and writes that shard to +all four peer output tensors. The next column-parallel projection waits for the +full `[8000, 5120]` normalized tensor before launching one exact CUTLASS GEMM. + +For exact `M=8000`, the candidate changes this boundary contract: + +1. Reduce-scatter and Gemma RMSNorm produce only the rank-owned + `[2000, 5120]` normalized shard, retaining the accepted FP16 reduction and + FP32 residual order. +2. The following QKV/QKVZ or gate/up projection consumes the local shard first + and then the three peer shards as they become visible. +3. Four `M=2000` CUTLASS subproblems write directly into their corresponding + row slices of the normal full projection output. The K reduction and output + row order are unchanged. +4. A replay-safe TP barrier protects peer visibility. A later iteration may + replace the four launches with a pointer-array/grouped persistent launch, + but only if the simple version proves that avoiding the materialized + all-gather is profitable on V100 NVLink. + +This is the SM70 analogue of an all-gather/matmul consumer pipeline. Generic +vLLM Sequence Parallelism and AsyncTP cannot be enabled unchanged here: + +- SM70 has no generic SP admission threshold; +- the accepted projection is a custom block-FP8 dequant + CUTLASS operation, + while the generic fusion pattern matches `aten.mm` (or newer scaled-mm + routes); +- the current exact-8K graph uses SM70-specific mixed-dtype residual semantics + that must not be replaced by a generic norm. + +The prototype therefore reuses the generic topology, but implements a bounded +SM70 custom consumer. It does not use Ampere/Hopper-only `cp.async`, TMA, or +native FP8 Tensor Core instructions. + +## Rejected paths that must not be repeated + +- Side-stream two-workspace FP8 weight dequantization across RMSNorm was + bitwise correct but regressed `7.98956 -> 8.00819 ms`. +- cuBLAS algorithm 106 lost to the accepted CUTLASS operation. +- Tile-ready down/output projection overlap failed to beat the full-model + baseline. Its acceptance runs were around 1.555-1.559 s pure prefill. +- CUDA graph wait rewrites regressed to roughly 1.565-1.572 s pure prefill; + earlier variants also errored or deadlocked. +- A prior generic symmetric-memory down+norm pipeline regressed its isolated + boundary. It is evidence against copying AsyncTP literally, not against the + new bounded consumer that removes the normalized all-gather. + +## Validation ladder and rollback + +1. Reproduce the accepted baseline on current main with the exact official + contract and record wall/pure-prefill samples, route hits, binary hash, and + GPU ownership. +2. Add an operator benchmark for fused norm plus the following exact QKV and + gate/up projections. Compare the complete boundary, not GEMM alone. +3. Require finite output and bitwise equality where the unchanged per-row + CUTLASS reduction permits it. Any numerical relaxation requires a separate + model-quality gate and cannot be accepted from a speed smoke. +4. Integrate only after the complete operator boundary beats the accepted + route by enough to plausibly save 107.5 ms over 127 boundaries. +5. Run a full candidate/control/candidate sandwich, seven official sampling + repeats, and a source-matching Nsight trace. Report request-wall and pure + prefill throughput independently. + +The candidate is exact-shape and default-off until acceptance. Its rollback +gate will be `VLLM_SM70_FP8_PREFILL_SHARD_CONSUMER=0`. The existing exact +CUTLASS rollback remains `VLLM_SM70_FP8_PREFILL_CUTLASS=0`. diff --git a/docs/design/sm70_qwen38_fp8_prefill_decay.md b/docs/design/sm70_qwen38_fp8_prefill_decay.md index 2ef310aaf6..b4e62eaf69 100644 --- a/docs/design/sm70_qwen38_fp8_prefill_decay.md +++ b/docs/design/sm70_qwen38_fp8_prefill_decay.md @@ -398,10 +398,139 @@ at most `1e-3` and relative L2 at most `1e-2`; greedy or bitwise identity is not required. The complete Flash-V100 policy suite passes 112 tests on a real V100, including a direct architecture-OOM-to-dense-fallback test. +## 2026-08-26 isolated TP4 256K boundary gate + +The first retained 261,888-token attempt did not establish a kernel or +sampling deadlock. A second TP4 job had started on the same GPUs without an +atomic claim. Its rank 2 worker, PID 827324, occupied 14.19 GiB on GPU 2 while +the prefill worker occupied 17.52 GiB. The prefill worker then failed a +134-MiB gate/up allocation with only 6.19 MiB free. Peer ranks waited for the +failed worker and the engine eventually reported `sample_tokens` RPC timeout. +The worker OOM is the primary failure; the RPC timeout is only the delayed +symptom. + +The corrected launcher refuses a direct unclaimed start, takes an atomic TP4 +GPU claim, and uses a case-local vLLM, TorchInductor, and Triton cache. During +the first corrected launch it rejected a real concurrent starter holding 368 +MiB on GPU 0, released the provisional claim, and retried only after all four +GPUs were idle. The measured contract is Qwen3.8-27B-FP8, TP4 V100-SXM2-32GB, +261,888 prompt tokens plus 32 output tokens, 262,144 maximum model length, +8,192 maximum batched tokens, one sequence, FP8 E5M2 KV, chunked prefill, +Mamba align mode, Flash-Attention-V100, CUDA graphs, and the model's official +thinking sampling (`temperature=1.0`, `top_p=0.95`, `top_k=20`, seed +20260825). The isolated profile exposes 18.64 GiB for KV and the single-request +CUDA graph occupies 0.20 GiB; the overlapped failure exposed only 7.11 GiB for +KV and captured 1.50 GiB of graphs under its old 16-sequence benchmark +contract. + +The first isolated A/B accidentally resolved the worktree extension symlink to +an earlier integration binary with SHA256 `63af4602...`. It is retained only +as diagnostic evidence. The final matched A/B explicitly loads the clean +production-patch binary with SHA256 +`91ad27ec8062edecb3f1a584e4db93cb4bce6cd11e43abf887039aec469c2314`; +the launcher now verifies that hash before engine startup. + +| Route | Pure prefill | Prompt throughput | Request wall | Output gate | +|---|---:|---:|---:|---| +| architecture disabled | 128.336222 s | 2040.64 tok/s | 129.263317 s | coherent 32-token text | +| 8K..256K architecture | 89.234794 s | 2934.82 tok/s | 90.167678 s | reject: 32 token IDs are zero | + +The architecture performance lane is 1.43819x faster, saves 39.101429 seconds, +and lowers pure-prefill time by 30.4680%. It is not an accepted model result: +the matched control emits coherent text with hash +`c3d3db079d1611c0462b3cf6cb55e83438123b3f607492b68ab2d6c32e114872`, +whereas the candidate emits 32 zero IDs with hash +`8a86dd30276a7f38fb6cf14439c812e4446227c5a48c7373fd5a9a7c07c4f97f`. +Until this model-level quality regression is fixed, 2040.64 prompt tok/s is +the usable 256K result and 2934.82 prompt tok/s is performance-only evidence. +The Draft remains blocked from promotion and the architecture rollback stays +`VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL=0`. + +Route evidence confirms that the candidate is not a final-chunk-only speedup. +Each rank reports 496 long-architecture hits: 31 full prefix-bearing chunks +times 16 full-attention layers. The first no-prefix chunk uses 16 exact +Split-D calls and the final partial chunk uses 16 tail-pad calls. The matched +control reports zero long-architecture hits. + +## 2026-08-26 half2 quality audit and stable 256K recovery + +The zero-token failure is not inherent to the split QK/PV and online-prefix +architecture. It is caused by the later raw-logit half2 shortcut. That route +skips row-max/statistics, fixes the prefix softmax peak at zero, and evaluates +the prefix weight as a degree-five FP16 Taylor polynomial of `score / 4`, +followed by two squarings. The approximation is narrow-domain: at scores +`-10`, `-12`, and `-16`, its weight is respectively about `15.95x`, +`2.89e4x`, and `1.38e9x` the exact exponential. Positive raw scores above the +FP16 exponential range overflow because no max is subtracted. + +The failure reproduces in the exact Q8000/KV40000 operator shape. At input +scale 1.25, 17 of 12,288,000 outputs are already non-finite. Scale 1.5 leaves +1,078,266 non-finite outputs, and scale 2.0 makes every output NaN. The +row-max/row-sum architecture remains fully finite at scale 2.0. This separates +a numerical-invalid shortcut from an architecture-level accuracy cost. + +An opt-in production score audit then sampled 1,048,576 scaled QK scores on +the first eligible 40K full-attention layer of every TP rank. The observed +score ranges were `[-9.28125, 14.41406]`, `[-8.74219, 7.22266]`, +`[-9.95312, 10.96875]`, and `[-9.10938, 18.25]`. After that first half2 +attention layer, Q, K, and sampled scores are NaN from audit index one onward +on all four ranks. The run emits eight zero token IDs. Thus the first +accelerated full-attention layer poisons later layers; the final zero-token +output is a downstream symptom. + +The quality recovery stops applying the raw-logit half2 patch and restores +the prior row-max/row-sum transform and online max-aware merge. Its stable +shape gate is extended from KV 128K to KV 256K without adding KV-proportional +workspace: the first 8K no-prefix chunk remains on exact Split-D, and every +full chunk from 16K through 256K uses the stable architecture. The rejected +half2 patch remains available in Git history but is removed from the source +tree so it cannot be applied accidentally. A future packed exponential must +operate only after max-shift and pass the adversarial finite-output gate. + +The rebuilt CUDA 12.8/SM70 binary has SHA256 +`f9f9acbc610c87fce9984e8fbd93fe0c8fa59887542123a74b3eaef6d3b8abf9`. +Its 256K operator result is: + +| Route | Latency | Useful causal TFLOP/s | Finite outputs | Relative L2 | Max abs diff | +|---|---:|---:|---:|---:|---:| +| exact Split-D | 289.605637 ms | 42.77 | 12,288,000 / 12,288,000 | reference | reference | +| stable max-aware architecture | 202.749443 ms | 61.09 | 12,288,000 / 12,288,000 | 0.007176 | 1.373e-4 | + +The model gate uses the same TP4, FP8 E5M2 KV, chunk-8192, Mamba-align, +CUDA-graph, and official-sampling contract as the clean 256K A/B. It runs a +40K sentinel followed by the 261,888-token contract in one engine: + +| Case | Pure prefill | Prompt throughput | Output gate | +|---|---:|---:|---| +| 40K stable sentinel | 9.535619 s | 4194.80 tok/s | coherent 32-token text, hash `426f6f...ca90` | +| 256K architecture disabled control | 128.336222 s | 2040.64 tok/s | coherent hash `c3d3db...14872` | +| 256K stable architecture | 107.380095 s | 2438.89 tok/s | exact same hash `c3d3db...14872` | + +The stable route lowers 256K pure-prefill time by 16.3291% and raises prompt +throughput by 19.5158% against the quality-valid control. Each rank reports +560 stable architecture hits across the two cases: 64 for the four eligible +40K chunks and 496 for all 31 eligible 256K chunks, each multiplied by 16 +full-attention layers. This supersedes the earlier quality rejection: the +output blocker is resolved while retaining a quality-valid 61-TFLOP/s +attention path. The updated Flash-V100 policy suite passes all 112 tests on a +claimed V100. The private current-main audit removes the rejected half2 source +patch and promotes this stable 16K-through-256K route with +`VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL=0` as the rollback. + ## Artifacts - D256 GQA architecture artifacts are retained outside Git; their private filesystem location is intentionally not committed. +- Final clean-binary 256K candidate/control JSON: + `model/results/tp4-256k-owned-official-sampling-clean91ad-{m080,control-m080}.json` + under the 80-TFLOP task root, with matching `.log` and `.claim.log` files. +- Claim-enforcing 256K runner: + `model/run_tp4_prefill_256k_owned.sh` under the same task root. +- Half2 audit and stable recovery evidence under the 80-TFLOP task root: + `audit/scale{1p25,1p5,2}-kv40000.json`, + `audit/exact-stat-{scale2-kv40000-new,kv256000}.json`, and + `model/results/tp4-exact-stat-output-fix-40k-256k-m080.json`, with the + matching model log. - Final operator A/B/A: `results/torch-architecture-scorecache-final-v1-aba.json` under that task root. diff --git a/docs/design/sm70_tile_runtime_exploration.md b/docs/design/sm70_tile_runtime_exploration.md index 0fe732e7a1..1a594ba074 100644 --- a/docs/design/sm70_tile_runtime_exploration.md +++ b/docs/design/sm70_tile_runtime_exploration.md @@ -19,12 +19,12 @@ decode path: Relevant public references: -- TileRT README: https://github.com/tile-ai/TileRT -- TileRT execution-model blog: https://www.tilert.ai/blog/speed-as-the-next-scaling-law.html -- TileRT 1000 TPS blog: https://www.tilert.ai/blog/breaking-1000-tps.html -- DistFuse tile-wise GEMM/all-reduce paper: https://mcanini.github.io/papers/distfuse.hotinfra24.pdf -- TokenWeave paper: https://arxiv.org/html/2505.11329v2 -- CUTLASS-native Distributed GEMM write-up: https://blog.shi-labs.com/distributed-gemm-88be6a481e2b +- TileRT README: +- TileRT execution-model blog: +- TileRT 1000 TPS blog: +- DistFuse tile-wise GEMM/all-reduce paper: +- TokenWeave paper: +- CUTLASS-native Distributed GEMM write-up: The TileRT repository currently ships Python wrappers and binary backends, not the core engine source. The actionable implementation detail therefore comes @@ -314,8 +314,8 @@ Correctness: `gated_silu` epilogue produced `max_abs=0`; interleaved normal output restored to the original order also produced `max_abs=0`. - End-to-end Qwen3.6-27B-AWQ TP2 token hashes matched baseline: - - 16-token smoke: `3d326cfe4c2bcd0aae49b11ed684dc2baff445be00e7eff6472cb014eb269cdb`. - - 128-token repeat-3: `25acc1257b65019cd7398d36105134378675d9de82bc7d30c6b49d78b1f7b755`. + - 16-token smoke: `3d326cfe4c2bcd0aae49b11ed684dc2baff445be00e7eff6472cb014eb269cdb`. + - 128-token repeat-3: `25acc1257b65019cd7398d36105134378675d9de82bc7d30c6b49d78b1f7b755`. Performance, 27B-AWQ TP2 on GPU0/1, `input_len=512`, `output_len=128`, `repeat=3`, Flash-V100 compile graph: @@ -390,9 +390,9 @@ Correctness: both repeats produced token hash `1ccba72c62311edee3c604b4d41cdeddad72adb96fd6666a1af609c59708fd68`. - Tail-worker reduce also matches the normal path: - - short smoke, `output_len=4`: token hash + - short smoke, `output_len=4`: token hash `e3b3ebc930b09ed8921c4e20f0c8dd6fb93f0bd51c382cdc33a485050ef60939`; - - stable comparison, `output_len=32`: both repeats produced token hash + - stable comparison, `output_len=32`: both repeats produced token hash `1ccba72c62311edee3c604b4d41cdeddad72adb96fd6666a1af609c59708fd68`. Performance, 27B-AWQ TP2 on GPU0/1, `input_len=512`, `output_len=32`, diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index bc1aa13bd3..b951142f2c 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -24,6 +24,112 @@ Goal: sampling, and normal MoE architecture wherever possible. - Do not carry unsafe quality-regressing routes into the main path. +## Qwen3.8 Flash Next NVFP4 prefill acceptance, 2026-08-27 + +- Review scope: Draft PR #23 on + `codex/v100-qwen38-flash-next-prefill-20260827-000848`, stacked on the + Qwen3.8 Flash Next bring-up branch. This scope is prefill-only; decode and + MTP acceptance remain separate work. +- Frozen route: four V100-SXM2-32GB GPUs, TP4/PP1, V2 engine, ModelOpt NVFP4, + FP16 activations and KV cache, MTP off, `FLASH_ATTN_V100`, FlashQLA-SM70 GDN + prefill, pinned-host PLE, prefix caching off, chunked prefill on, and + `VLLM_COMPILE` with PIECEWISE CUDA graphs. +- The matched 8192-token candidate has mean pure prefill time `2.315786s`, or + `3537.46 tok/s`, and mean TTFT `2.322317s`. The earlier same-route baseline + has mean pure prefill time `21.339760s`, or `383.88 tok/s`: a `9.2149x` + speedup and `89.15%` latency reduction. The baseline generated 512 tokens + and used FULL_AND_PIECEWISE while the candidate generated 16 tokens and used + PIECEWISE; only request-metric prefill time is compared, and prefill uses the + PIECEWISE dispatch in both runs. +- The accepted change admits the two real 512-expert Flash Next prefill GEMM + contracts to grouped TurboMind NVFP4 execution, uses four QSA sparse-kernel + warps for exact SM70 prefill shapes, and permits an absolute prebuilt + FlashQLA-SM70 extension path so all TP ranks avoid redundant JIT builds. + Focused validation is `15 passed, 5 skipped`; exact QSA probes are bitwise + equal and improve the 8192-token kernel profile from `1320.204ms` to + `85.261ms`. +- Long-context pure-prefill results from one guarded model startup are: + 32768 tokens `10.740389s` / `3050.91 tok/s`, 65536 tokens `25.020538s` / + `2619.29 tok/s`, and 131000 tokens `63.960919s` / `2048.13 tok/s`. The 32K + and 64K repeat token hashes are stable, and the arithmetic and Chinese + quality outputs match the accepted 8K run exactly. +- A follow-up SM70-only QSA tile reduction from 64 to 32 columns passed the + matched long-context gate. Pure-prefill time is `9.044346s` at 32K, + `21.517608s` at 64K, and `56.609353s` at 131K, corresponding to + `3623.04`, `3045.69`, and `2314.11 tok/s`. Relative to the preceding + candidate this is a `1.1875x`, `1.1628x`, and `1.1299x` speedup, with + latency reductions of `15.79%`, `14.00%`, and `11.49%` respectively. All + five quality/performance case token hashes remain exactly equal to the + matched baseline. +- Exact-shape QSA microbenchmarks explain the end-to-end change. At 8192 rows, + the 32-column tile measures `56.723ms` versus `92.730ms` for 64 columns + (`1.6348x`); at 512 rows with four splits it measures `4.397ms` versus + `6.385ms` (`1.4520x`). Both have maximum FP16 difference `0.0001220703125` + and no non-finite values. A 128-column tile is rejected before execution + because its `139264`-byte shared-memory requirement exceeds V100's + `98304`-byte limit. +- The next hotspot audit inspected the cached SM70 PTX instead of screening + more launch shapes blindly. Both QSA scoring kernels contain no `mma`/`wmma` + instructions: selected attention compiles to roughly 1026 static FP32 FMA + instructions, and the four-head D128 index scorer compiles to roughly 2048. + The latter grows with all visible compressed blocks and explains most of the + additional 131K decay. +- A D256 selected-attention WMMA prototype is rejected. It has real SM70 HMMA, + no local-memory spill, and passes the resource gate, but measures only + `51.808ms` versus `55.682ms` at 8192 rows (`1.0748x`) and its output relative + L2 is `0.5627`. At 512 rows it is only `1.2532x`. It is neither correct nor + fast enough to justify a production integration or model restart. +- The retained long-indexer candidate gathers one request's paged FP16 MQA + keys once, applies a cuBLAS FP16-input/FP32-output Tensor Core GEMM in a + bounded workspace, and fuses per-head ReLU, head sum, visibility and scale + into one Triton epilogue. At 1024 rows and 131K context, scoring including + paged gather moves from `35.767ms` to `2.108ms` (`16.964x`), with maximum + absolute error `8.11e-6` and relative L2 `2.90e-7`. The complete + score/top-k/expand path moves from `38.792ms` to `3.118ms` (`12.441x`). + Random stress selection sets are exact for 1023/1024 rows; the sole mismatch + is a boundary near-tie, and final sparse-attention cosine is `0.9999992`. + Production admission is limited to exact SM70, FP16 `[rows,4,128]`, FP16 + `[pages,page,1,128]`, one request, at least 512 rows, and at least 1,048,576 + row-column score elements; decode, short work and generic batches retain the + paged Triton route. The measured crossover is `0.733x` at 1K context, + `2.164x` at 2K, `4.900x` at 8K, `8.064x` at 16K and `11.102x` at 32K. A real + TP4 token-hash and quality gate remains required before this candidate is + accepted or reported as an endpoint improvement. +- Raising the NVFP4 MoE tuning limit is rejected independently of the model + gate. Real layer-0 weights and the exact 8192-token/top-10 routing shape + measure default W13/W2 at `4.320/2.382ms` (`31.07/28.18 TFLOP/s`) and the + measured policy at `4.343/2.381ms`; combined tuning is `0.9966x`. The real + router distribution is highly skewed (0/58/160/2150 minimum/median/mean/ + maximum rows per expert), but TurboMind's measured dispatch does not improve + it. Do not spend another startup autotuning this same shape. +- No new GDN candidate is claimed by this follow-up. A cold standalone 8192 + token original-TileLang probe under `/usr` CUDA attempted to rebuild the + TileLang cumsum kernel and failed in the installed TileLang BF16 fallback + headers. The guarded full-model run reused the accepted production cache and + route successfully. Do not treat the standalone failure as a GDN runtime + regression or repeat it without the source-matched TileLang toolchain/cache. +- At `gpu_memory_utilization=0.85` and `max_model_len=140000`, vLLM reports + `3.09 GiB` available KV cache, `250028` token capacity, and `1.79x` 140K + concurrency. Measured peak device use is `29090 MiB` (`28.408 GiB`) per + rank; minimum host `MemAvailable` is `47.912 GiB`, with `247.921 GiB` swap + still free. +- Do not repeat the rejected `gpu_memory_utilization=0.82` 140K startup. It + fails cleanly before requests because `1.08 GiB` KV cache is below the + required `1.73 GiB`, with an estimated maximum length of 86800 tokens. The + single `0.85` retry passed the capacity gate and all planned 32K/64K/131K + cases, so no third startup is needed. +- Retained local evidence is under + `.artifacts/qwen38_flash_next_prefill_20260827/results/`, notably + `grouped-v1-piecewise-8192x16.json` and + `grouped-v1-piecewise-longctx-u085.json`. The QSA follow-up is retained as + `grouped-v2-qsa-bn32-longctx-u085.json`, with microbenchmarks in + `qsa-blockn-sweep-{8192,512-split4}.json`. The corresponding PR evidence is + recorded in the + [8K result](https://github.com/yangzhuxinyzx/1Cat-vLLM-private/pull/23#issuecomment-5436515073) + and + [long-context result](https://github.com/yangzhuxinyzx/1Cat-vLLM-private/pull/23#issuecomment-5437096030) + comments. + ## Active MRV2 DFlash2 campaign, 2026-08-20 - Integration base: `onecat/main@7aede2cf010d92815c9d7bff25867b4fa009b6cb`. @@ -42336,6 +42442,74 @@ Interpretation: separate from accepted kernel repairs. Small policy substitutions still need their own current-source static/argument-equivalence checks before merge. +## 2026-08-27 GLM-5.3-Flash NVFP4 SM70 bring-up + +- Pinned `LibertAIDAI/GLM-5.3-Flash-NVFP4` at revision + `9e0d74e3cef17f634e84fb8e2223707e02616290`. The acceptance topology is eight + V100s arranged as TP4/PP2, FP16 activations, no MTP, batch one, exact + 1024-input/256-output measurement, and at least 70 steady pure-decode tok/s. +- Audited vLLM PR 53906 and SGLang PRs 36507/36513. Their GLM architecture and + scheduler semantics were used as references, but their Hopper/Blackwell FP8 + validation is not an SM70 performance baseline. The conflicting broad vLLM + PR was ported by required surface instead of merged wholesale. +- Added `glm5_next`, KDA bounded-gate support, mHC PP state, sparse MLA/DSA + indexer and KPool ownership, hybrid-cache grouping, GLM processor/config, + stacked weight mapping, and exact-SM70 ModelOpt NVFP4 MoE admission. +- V100 cannot perform the KPool writer's native FP8 conversion. Exact SM70 now + emits software E4M3 bytes through `uint8` pointers and performs FWHT query + rotation plus sparse scoring in FP16 through the established SM70 HMMA + indexer. The 28-test GPU KPool suite is byte-exact against the reference. +- A full local source build completed and registered the NVFP4 SM70 and MoE + custom operators. Focused CPU integration reports 64 passing tests. An exact + eight-rank TP4/PP2 `meta` construction smoke reports stage sizes 23/22, + selects `ModelOptNvFp4SM70MoEMethod` on all ranks, and preserves all four mHC + PP tensors. +- These are route and construction results, not an accepted model result. Keep + the real checkpoint load, three-request output identity, sampling quality, + CUDA Graph route hit, and matched pure-decode benchmark open until an + uncontended eight-GPU run supplies retained logs and JSON. +- The checkpoint download is now complete: all 120 shards are present at the + pinned revision (about 181 GiB). Exact eight-V100 TP4/PP2 dummy-weight eager + requests pass with both FP16 KV and `fp8_e4m3` KV. The FP16 run reports + 66,355 cache tokens; packed E4M3FN reports 111,957 tokens (about 1.69x). +- Added a GLM-specific 520-byte latent-KV format: 512 software-E4M3FN data + bytes plus eight UE8M0 scales per token. `_sm70_glm5_fp8_kv_insert_kernel` + writes it. B1 decode keeps the cache packed, gathers/dequantizes the selected + rows into a fixed-width FP16 workspace, then uses two Tensor Core GEMMs. This + replaced the direct scalar sparse kernel on the accepted decode route. +- Fixed the upstream GLM KDA state contract to match its merged Q/K/V + convolution implementation. Each layer now exposes two states and two copy + functions: merged FP16 convolution plus FP32 recurrent state. The original + four-state declaration failed at the first real model call after cache + allocation. +- Replaced the exact GLM mHC SM70 PyTorch fallbacks. M=1 fused decode uses the + existing DeepSeek-V4-style FP32 staging path plus a native CUDA + Sinkhorn/residual-mix/RMSNorm final stage; + standalone pre/post and M>16 fused prefill use dedicated SM70 Triton + reductions/post mapping. A direct M=17 probe reproduces the rejected generic + TileLang path's 33 SM70 BF16-header compile errors. The replacement passes + 17 focused V100 tests and a full input-17/output-2 TP4/PP2 FP8-KV request. + Retained route log: + `/data/minimax-h3/task-cache/glm53-nvfp4-sm70-20260827/dummy_fp8_tp4pp2_i17_o2_fast_mhc.log`. +- Layer-zero mHC now pre-sums the four identical initial-stream weights. A + V100 M=1 microbenchmark including the required residual expansion measures + about 0.155 ms for the expanded path and 0.140 ms for the broadcast path. + This is a local kernel result, not an end-to-end speed claim. +- The first real no-compile CUDA Graph baseline measured 5.35 steady decode + tok/s. Replacing scalar packed-FP8 sparse attention with bounded gather/ + dequant plus Tensor Core GEMMs raised the retained short benchmark to + 45.83 tok/s (21.82 ms TPOT), with the same 16-token hash + `70007811d61c68bb6ec6b4ac5758744f2d4c6b64b51b6b1352f380091c990902`. + Nsight Systems reports 95.09% graph-node coverage and no remaining sparse + attention hotspot. This 16-input/16-output result is an optimization + checkpoint, not the required 1024/256 acceptance result. +- Added a graph-safe exact-SM70 TP4/B1 CUDA kernel that fuses KDA's two local + `128 -> 2048` f/g projections. It measures 6.28 us versus 23.02 us for two + cuBLAS launches (3.66x) and preserves FP32 accumulation with FP16 output. + Focused mHC/KDA graph and numerical coverage passes 20 tests with one skipped; + the next real checkpoint run must confirm its end-to-end gain and unchanged + greedy token hash. + ## 2026-08-22 Qwen3.6-35B-A3B ModelOpt NVFP4 TurboMind route - The checkpoint at `/data/models/Qwen3.6-35B-A3B-NVFP4` is ModelOpt 0.37.0 @@ -43449,16 +43623,35 @@ Interpretation: the current direct-top6 pipeline from `54.414` to `46.519 us/layer`, a `0.339 ms/token` 43-layer service projection. The reproducible JSON is under `/data/models/v100-dsv4-0731-pp2tp4-mxfp4-direct-order-screen-20260826-r1/`. - This is an operator gate only; - the opt-in `VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE` route still requires a - matched PP2/TP4 endpoint and aggregate-quality gate. + The exact B1, replicated-expert, top-k-6 route is therefore default-on after + source audit. `VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE=0`, + `VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE=0`, and + `VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE=0` remain explicit rollbacks; other + token counts, expert mappings, top-k shapes, dtypes, and devices retain the + existing path. This is operator performance evidence, not a full endpoint + throughput claim. - The first matched endpoint reached `69.949 token/s`, but its quality gate - failed with incoherent output. The direct-order branch had reused - `compact_expert_offsets` as if it still contained `[0, 1, ..., 6]`; M2--M8 - warmup and verifier compaction can overwrite that workspace before B1 graph - capture. The corrected branch uses the immutable `slot_expert_offsets` - buffer. Keep the route opt-in until a post-fix full-model and GSM8K-64 gate - pass; the failed endpoint is rejection evidence, not a speed baseline. + failed with incoherent output and is rejected as a speed baseline. M2--M8 + warmup and verifier compaction can overwrite `compact_expert_offsets` before + the B1 graph is captured, while direct-order incorrectly assumed that buffer + still contained `[0, 1, ..., 6]`. Direct-order now reads the immutable + `slot_expert_offsets` buffer. This restores the proven one-row-per-route + contract without changing the kernel, route order, or default-on/rollback + policy; a poisoned-workspace regression covers the failure mode. +- The exact TP4 QNorm/RoPE plus FP8-KV insertion route combines the existing + 16 Q-head programs and eight packed-KV writers into one graph node without + changing their arithmetic. Across 64 changing-input CUDA Graph patterns, + both transformed Q and the complete cache storage are bitwise equal. Median + operator time falls from `9.974` to `4.076 us` (`2.45x`), projecting + `0.254 ms/token` across 43 layers. Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-qnorm-kv-fusion-screen-20260826-r1/`; + `VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4=0` is the rollback. +- Sparse MLA QK d-split now uses one 16-head group for the exact TP4 shape. + V100 medians improve by about 43% at C4 and 30--31% at C128/SWA. C128 and + SWA are bitwise equal; the C4 and changing-input comparisons remain finite + with at most one changed FP16 element per pattern, `1.526e-5` max absolute + difference, and `1.807e-5` max relative L2. Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-sparse-head16-screen-20260826-r1/`. ## 2026-08-25 DFlash2 n-gram hybrid @@ -43504,30 +43697,29 @@ Interpretation: acceptance. The assistant remains explicit opt-in until matched evidence shows that full-query skips repay host lookup overhead for a workload. -## 2026-08-26 PR #283 aggregate-quality acceptance +## 2026-08-26 PR #283 aggregate-quality decision (superseded) -- The numerical policy does not require greedy or token-stream identity. - Candidate outputs must remain finite, satisfy dtype-appropriate operator - error bounds, and preserve matched aggregate task quality. Per-example - regressions and improvements remain visible diagnostics rather than an - automatic identity gate. +- This section records the historical aggregate-only decision. It is + superseded by the no-per-sample-regression policy below: an aggregate tie or + improvement cannot compensate for any control-correct sample becoming + wrong, and an algebraically exact candidate must preserve the matched output + stream. - The full serialized-FP8 QPN8 candidate improves the matched PP2 x TP4 B1 no-spec endpoint from 59.248 to 64.359 token/s (`+8.63%`) and reduces mean TPOT from 16.878 to 15.538 ms. Its paired GSM8K-64 result remains 63/64 with zero invalid answers: one baseline-correct answer regresses and one - baseline-wrong answer improves, so aggregate quality does not decline. + baseline-wrong answer improves. Under the current policy, the + baseline-correct-to-candidate-wrong transition rejects this result. - Operator checks remain finite with relative L2 at most `6.05e-4`, cosine at least `0.9999997`, and maximum absolute difference at most `0.00390625` in the recorded screen. All 365 audited block scales are finite, positive powers of two and survive the layout scale transform exactly. These are numerical bounds, not a claim of bitwise equivalence. -- QPN8 is therefore default-on only for its validated engine contract: exact - SM70 block-FP8 operator roles/shapes/layouts, PP2 x TP4, one sequence, no - DBO or explicit ubatching, and no speculative decoding. Missing operators - or workspace allocation fall back to TurboMind. Either - `VLLM_SM70_FP8_QPN8=0` or `VLLM_SM70_FP8_QPN8_PP2_TP4=0` is an explicit - rollback. Admission never reads model name, checkpoint, `model_type`, or - architecture identity. +- QPN8 is now default-off for the PP2 x TP4 contract. Explicit + `VLLM_SM70_FP8_QPN8_PP2_TP4=1` remains available only for controlled + diagnosis and still requires its exact tensor/runtime contract. The matched + model-level regressions below override the earlier operator-error and + aggregate-score rationale. - Metadata-free PP transfer is also default-on for its exact SM70 B1 schema: PP2 x TP4, FP16 `[1,4,4096]` contiguous replicated hidden state, CUDA Graph, no sequence parallelism, no DBO/ubatching/speculation, and one sequence. @@ -43536,32 +43728,50 @@ Interpretation: original metadata plus TP reconstruction path; an admitted-but-invalid runtime tensor fails fast before communication. The rollback is `VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER=0`. -- The paired quality tools now pin dataset hashes and the complete evaluation - contract, validate artifact self-consistency, and gate aggregate HumanEval, - LongBench, GSM8K, and needle-retrieval quality. They report directional - sample flips but do not treat greedy identity as task quality. -- A post-merge real-weight accuracy audit replaces the PP2 x TP4 route's - fastest single-pattern schedules with four-pattern accuracy-first schedules. - Fused WQA/WKV, attention WQ-B/WO-B, grouped WO-A, and shared-expert down - retain at least `99.707%` exact FP16 elements and at most `1.444e-5` - relative L2 error against same-process TurboMind. Shared-expert gate/up - retains only `53.516%` exact elements with its fused activation and is - removed from the default QPN8 tensor contract; the replicated indexer WQ-B - remains excluded. -- The latest-main production source gate covers M=1 and M=9, changing-input - CUDA Graph replay, FP32 and TurboMind references, grouped WO-A layout, and - both excluded roles. All five admitted projections pass, both excluded roles - remain on TurboMind, and the route reuses one 16-MiB workspace. Evidence is - under `/data/models/v100-dsv4-0731-pp2tp4-qpn8-accuracy-source-gate-20260826-r2-main/`. - The isolated weighted projection is `6.067 -> 3.321 ms/token` of stage-sum - FP8 service; it is not an endpoint TPOT claim. +- The paired quality tools pin dataset hashes and the complete evaluation + contract, validate artifact self-consistency, and gate HumanEval, LongBench, + GSM8K, and needle-retrieval quality per sample. A control-pass to + candidate-fail transition is an automatic rejection. +- A post-merge four-pattern, real-weight matrix separates a single numerical + outlier from five already-bounded fast schedules. Fused WQA/WKV, attention + WQ-B/WO-B, grouped WO-A, and shared-expert down retain at least `99.707%` + exact FP16 elements and at most `1.641e-5` relative L2 against same-process + TurboMind with their existing speed winners. Their accuracy-first schedules + are about 15--36% slower at the isolated operator while only reducing errors + that are already in the `1e-5` range, so the production route keeps the + speed winners. The matrix is under + `/data/models/v100-dsv4-0731-pp2tp4-qpn8-accuracy-matrix-20260826-r1/`. +- Shared-expert gate/up is the material outlier: its fused activation retains + only `53.516%` exact elements and reaches `5.989e-4` relative L2. It is the + only role removed from the default QPN8 tensor contract and now falls back + to TurboMind. The replicated indexer WQ-B remains excluded. This narrows an + existing hardware/topology/tensor route; it does not add a model-identity + admission rule or disable QPN8 globally. +- The conservative all-accuracy source gate covers M=1 and M=9, + changing-input CUDA Graph replay, FP32 and TurboMind references, grouped + WO-A layout, both excluded roles, and one reused 16-MiB workspace. All five + admitted projections pass and both excluded roles remain on TurboMind. + Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-qpn8-accuracy-source-gate-20260826-r2-main/`. + Its isolated weighted projection is `6.067 -> 3.321 ms/token` of stage-sum + FP8 service; this is a conservative lower-speed variant, not an endpoint + TPOT claim for the retained speed-winner schedules. - Withdraw the earlier causal attribution of a `63/64 -> 2/64` GSM8K result to fused-WQA QPN8 alone. That recorded endpoint also pinned five static FP8 tactics and enabled grouped WO-A; the static-FP8-only endpoint reproduces the same class of incoherent output. Those artifacts prove that the bundled - configuration fails, not which operator caused it. Keep the accuracy-first - schedule change unmerged until a clean matched pure-decode run and pinned - dataset gates close. + configuration fails, not which operator caused it. +- Clean matched PP2 x TP4, B1, no-spec runs now close the conservative + all-accuracy variant. QPN8-off records a median `63.6434 tok/s` and + `15.71254 ms/token`; the all-accuracy candidate records `63.9120 tok/s` and + `15.64652 ms/token`, about `+0.42%` and `-0.066 ms/token`. The paired pinned + GSM8K64 run is `62/64` versus `61/64`, with zero invalid outputs and one + directional difference at item 37 (`2 -> 0`) and no candidate win. This is + a model-level quality regression, so the candidate is rejected and the + broader PP2 x TP4 QPN8 route defaults off. Full artifacts are under + `/data/models/v100-dsv4-0731-pp2tp4-qpn8-accuracy-fullmodel-control-20260826-r1/` + and + `/data/models/v100-dsv4-0731-pp2tp4-qpn8-accuracy-fullmodel-candidate-20260826-r2/`. ## 2026-08-26 generic multimodal-tower UVA offload @@ -43608,3 +43818,357 @@ Interpretation: or eager call retains the existing pull collective. Explicit `=0` is the rollback. No model, checkpoint, `model_type`, or architecture identity is consulted. + +## 2026-08-26 historical PR #104 Gemma4 tower-memory audit + +- Historical PR #104 is not replayed as a branch. Five of its six commits are + already superseded by newer main-line attention, compile-warmup, MTP, and + metadata paths. Its remaining Gemma4 multimodal memory fix is rebuilt from + `origin/main@1492e6985a741da09329127b4d4171882b70402b` on + `codex/v100-gemma4-sm70-mm-sdpa-20260826-013904`. +- The old post-load `tower.to(dtype)` approach is incompatible with the + generic UVA tower offload merged in PR #297: it can replace the registered + UVA storage with a new device tensor. The current repair instead passes the + configured FP16 dtype to Transformers while constructing the vision/audio + towers, before tower registration and weight loading. This path is admitted + only on exact SM70 with a configured FP16 model dtype; other platforms and + dtypes keep their current construction and mask behavior. +- Both the current image-bucket encoder and the newer video-frame encoder use + a broadcast additive key-padding mask of shape `[B,1,1,N]`. Transformers + 5.5.3 returns this prepared 4D mask as-is, avoiding its normal expansion of + a 2D validity mask to `[B,1,N,N]`. The implementation constructs the mask + entirely on device and does not branch on a CUDA scalar. +- On a Tesla V100-SXM2-32GB, forced memory-efficient SDPA accepts the FP16 + broadcast mask and produces bitwise-identical output to the bool validity + mask in the focused probe (`max_abs=0`). The same forced backend rejects + BF16 on SM70 with `No available kernel`, confirming the dtype dispatch + defect rather than assuming it from model identity. +- At the production failure shape `B=2,N=10080`, Transformers 5.5.3 expands + the 2D mask to 203,212,800 bool elements (about 194.2 MiB). The broadcast + form stores 20,160 FP16 elements (about 0.04 MiB), a 10,080x element-count + reduction. Seven warmed CUDA-event samples measure medians of 1.346 ms and + 0.042 ms respectively; this is mask-construction evidence, not a full-model + throughput claim. +- CPU SDPA probes under Transformers 5.5.3 also produce exact equality for + padded bool/additive masks and for an all-zero additive mask versus no mask. + A tiny Gemma4 vision tower confirms that `AutoModel.from_config(dtype=fp16)` + creates FP16 parameters before loading. Ruff, formatting, `diff --check`, + and compile checks pass. The repository pytest entrypoint is currently + blocked before collection by the environment's unrelated stale + `mistral_common` (`NamedToolChoice` import); no service or end-to-end run is + claimed. + +## 2026-08-26 exact SM70 FP8 prescaled-M1 decode screen + +- The SM70 E4M3 transform normally applies the exact exponent-bias factor and + the checkpoint block scale as two FP16 multiplies. DeepSeek V4 UE8M0 block + scales are powers of two, so multiplying each prepared scale by 256 once at + load time removes the per-fragment exponent-bias multiply without changing + the MMA layout, split count, swizzle, FP32 accumulation order, or FP16 + epilogue. +- A same-process CUDA Graph A/B/B/A screen uses real layer-0 TP4-rank-0 + weights and the five non-gated M=1 decode shapes. All 320 changing-input + comparisons are bitwise equal. The fused WQA/WKV projection improves from + 63.330 to 20.577 microseconds (`3.078x`); the weighted five-shape projection + saves 1.867 ms/token warm and 1.821 ms/token after a 256-MiB cache scrub. + Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-fp8-prescaled-decode-screen-20260826-r1/`. +- Production admission is initially narrower than the operator screen: SM70, + UE8M0 128x128 block scales, PP2 x TP4, one sequence, no DBO/ubatching or + speculation, M=1, and the exact replicated K4096/N1536 fused-WQA/WKV role. + Model loading also requires the actual FP16 scales to remain finite and + bitwise reversible after the 256x exponent shift. Prefill, gated activation, + other shapes, other scale formats, missing operators, and unsupported scale + ranges retain the original TurboMind transform. The route is default-on for + the exact tensor/runtime contract; `VLLM_SM70_FP8_PRESCALED_M1_DECODE=0` + is the rollback. An explicit `=1` fails closed when its contract cannot be + honored. Admission never reads model name, checkpoint, `model_type`, or + architecture identity. + +## 2026-08-26 DeepSeek V4 packed-FP13 GEMV screen + +- The fixed batch-one SM70 auxiliary/router GEMV reads FP16 parameters that + originate from BF16 checkpoint weights. Packing the upper 13 FP16 bits + stores 32 values in 13 uint32 words and reduces weight traffic by 18.75%. + A full checkpoint scan covered all 188 eligible tensors and 354,680,832 + values. All 409,653 nonzero discarded tails were FP16 subnormals; no normal + value had a discarded bit and no nonfinite value was present. The worst + tensor-relative weight L2 is `3.19e-6`, with maximum absolute error + `4.18e-7`. +- Offline SM70 compilation of the single-reduction kernel reports 29 registers, + 17 global loads, 32 FFMA, eight shuffles, two barriers, and 227 static SASS + instructions. The exact source-BF13 variant uses 30 registers and 334 + instructions; rounded BF12 uses 31 registers and 327 instructions. FP13 is + therefore the only retained packed-GEMV candidate for production screening. +- The existing Nsight trace shows why the full 2.554-ms stage-sum GEMV service + must not be projected as endpoint savings. At a representative C4 layer the + default-stream FP8 projection takes 43.008 us while concurrent N2048/N512/N64 + GEMVs take 37.600/25.120/22.496 us. The auxiliaries complete before the + downstream dependency. Any wall benefit comes from lower concurrent HBM + contention plus the default-stream router GEMV, so an exclusive-V100 overlap + screen and then a matched endpoint are required. +- The production FP13 operator screen on real weights saves a call-count-weighted + `0.2074 ms/token` warm and `0.1998 ms/token` cold versus the FP16 operator. + An archived four-stream CUDA Graph screen at the source head was stable under + A/B/B/A ordering and saved `0.0550 ms` warm and `0.0310 ms` cold per C4 layer, + with 64/64 bitwise main-path results and maximum auxiliary relative L2 + `4.03e-7`. That overlap result is directional after the mainline FP8 M=1 API + split and is not relabeled as a current-main endpoint result. +- FP13 remains default-on under its narrow tensor/runtime contract because the + real-weight operator and graph screens show a bounded numerical delta and a + repeatable latency reduction, with no recorded task-quality regression. + `VLLM_SM70_DSV4_FP13_GEMV=0` is the explicit rollback. The route is not + reported as an endpoint gain until a matched model-level dataset A/B closes. + Admission still checks SM70, CUDA device equality, dtype, shape, layout, + output contract, and the actual weight bits; incompatible tensors retain + FP16. Packing occurs once after loading. The retained buffers add 549.66 MiB + without pipeline parallelism, or at most 282.34 MiB per rank for the audited + PP2 partition. + +## 2026-08-26 DeepSeek V4 PP2 x TP4 100-token/s continuation + +- The active target is single-request pure decode on eight V100s, PP2 x TP4, + CUDA Graph enabled, no speculative decoding and no DSpark. The acceptance + target remains `100 token/s` (`10 ms/token`) with matched model, quantization, + GPU/rank layout, input/output lengths, sampling, attention, KV-cache and + graph state. Prefill and TTFT are reported separately. The pinned quality + baseline remains coherent and scores GSM8K `63/64`. A narrowly bounded + operator may default on after graph, numerical and source gates pass, with + an explicit rollback; it is not reported as an endpoint gain until a matched + full-model benchmark exists. +- The source-`59adf8d89a` measurement control explicitly disables all FP8 QPN8 + routes and packed FP13 to isolate the remaining stack; that measurement + choice does not change the guarded FP13 runtime default. The control retains + the exact fused-WQA prescale, exact MXFP4 QPN M1, static PP transfer, TP4 push + all-reduce, fused Q/KV normalization, direct-order MoE, and exponent fold. It + scores GSM8K-64 + `64/64` with zero invalid answers, HumanEval `29/32`, and LongBench `44.740`. + Three matched 1024+256 pure-decode runs measure `73.534`, `73.542`, and + `73.539 token/s` (median `73.539`, `13.598 ms/token`). This is the accepted + no-DSpark baseline for the next trace; evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-private-quality-safe-prescaled-shared-gate-control-fullmodel-20260826-r2/`. +- The latest usable endpoint before the current queued stack is source + `089feeae7241daaaf7ccb2a4fb2a821925be5350`: QPN8-off records + `68.279 token/s`, and QPN8-on records a three-repeat median + `69.287 token/s` (`14.433 ms/token`). The corresponding coherent aggregate + quality observations are retained as controls. The earlier direct-order + `69.949 token/s` run is rejected because mutable compact offsets produced + incoherent output; the repaired route reads immutable slot offsets. +- The first exact combined stack at source + `6abbed64a95dd2e27e8a1daeb10118db926d857b` now closes its matched + endpoint and aggregate-quality run. Direct top-6/order, QPN8 including + WQA, static PP transfer, TP4 push all-reduce, fused Q/KV normalization, + MXFP4 exponent fold and joined FP16 auxiliaries are enabled; speculation, + DSpark and FP8 prescaling are disabled. Three 1024+256 pure-decode repeats + measure `72.372`, `72.353` and `72.381 token/s`, with median TPOT + `13.817 ms` and population standard deviation `0.0117 token/s`. All three + output-token hashes are identical. The Chinese smoke is coherent and + GSM8K-64 scores `62/64` with zero invalid answers, matching the preceding + accepted exact-stack observation and remaining one item below the pinned + `63/64` baseline. Under the current no-regression policy this is historical + speed evidence, not a quality-admissible endpoint. Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-qpn8-exact-combined-fullmodel-20260826-r1/`. +- The latest 26-step Nsight Systems trace measures a `16.535 ms` wall replay + interval. Excluding PP dependency waits, stage-sum service is led by FP8 + dense GEMMs `5.814 ms`, MXFP4 MoE `2.629 ms`, FP16 GEMV/compressor + `2.554 ms`, mHC `1.693 ms`, routing/activation `1.665 ms`, Q/KV work + `1.190 ms`, sparse MLA `1.103 ms`, TP all-reduce `0.745 ms`, and + LM-head/sample `0.415 ms`. The PP send/feedback rows are synchronization + residency, not additive transfer cost; actual payload kernels are only tens + of microseconds. Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-exact-stack-nsys-20260826-r1/`. +- A no-copy MXFP4 QPN prototype consumes the existing TurboMind E2M1/UE8M0 + pack directly. On real layer-0 TP4-rank-0 weights, W13 split-16/one-chain + plus W2 split-8/one-chain moves the W13--SwiGLU--W2 pipeline from + `51.299` to `37.934 us`, is CUDA-Graph stable, and is bitwise for the timing + input. Its 40-layer service projection is `0.535 ms/token`. A full-model + route-hit candidate measures `73.645`, `73.613` and `73.646 token/s`, with + median TPOT `13.579 ms`, identical output-token hashes, coherent Chinese + output and GSM8K-64 `62/64` with zero invalid answers. It is not + quality-admissible and is also not an isolated no-copy A/B: relative to the + `72.372 token/s` endpoint it also swaps FP8 QPN8 for the prescaled FP8 path, + so the observed `0.239 ms/token` difference cannot be assigned to MXFP4. + Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-mxfp4-tm-qpn-exact-combined-fullmodel-20260826-r1/`. +- The same no-copy kernel now defaults on under the exact B1 direct-order TP4 + six-route W13/W2 tensor contract; all other calls keep the existing + TurboMind dense-stage path. `VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE=0` is the + rollback. A loaded older extension without the new operator falls back when + the default is implicit, while explicit `=1` fails closed. A clean SM70 AOT + build produces only `sm_70` cubins with extension SHA256 + `9781c406cbd3cb38261df7be721bd8a7325e1f4cace6a06b6bb63014f5809507`. + The split-8 and split-16 kernels each use 56 registers and respectively + 1 KiB and 2 KiB shared memory. Worktree-source/new-extension binding and + the focused MXFP4 plus FP8-QPN compatibility suites pass (`41 passed`). + The selected operator pipeline is bitwise and graph-stable; no unmatched + operator projection is relabeled as a full-model speed result. The merged + private source now also has a clean CUDA 12.8 AOT build under + `/data/models/private-v100-dsv4-pp2tp4-defaults-build-20260826-r1/`. + Its extension SHA256 is + `ef6db06d54419029da87bbe159a36b0f5f0b68658797aceab69491b3f4421fa7`; + all 38 cubins are `sm_70`, and the split-8/split-16 kernels retain exactly + 56 registers with 1 KiB/2 KiB shared memory. A fresh interpreter binds the + private follow-up worktree to this `_C`, retains the installed `_moe_C`, + and sees both the MXFP4-QPN and FP8-QPN operators. With CUDA hidden, the + focused MXFP4 plus PP2-TP4 FP8-QPN suites pass (`39 passed, 4 skipped`). + This closes the merged-source AOT and host-policy follow-up. A merged-source + GPU route-hit and full-model quality endpoint remain follow-up evidence, not + a retroactive prerequisite for the already merged change. +- The exact FP16 auxiliary candidate joins the C4 `N=2048`, `N=512` and + `N=64` projections into one `N=2624` launch while preserving their original + FP32 FMA and reduction order. All joined outputs and the concurrent main + projection are bitwise across 64 changing inputs. Relative to the old FP16 + three-stream path, warm four-stream overlap moves `103.809` to + `51.639 us/C4 layer`; a real-V100 graph test, driver-free SM70 compile gate + and the combined endpoint quality gate pass. It is not ported to the newer + private source because that source already defaults to packed-FP13: the + three-stream FP13 path measures `50.879 us` warm and `70.072 us` cold, + versus `51.639 us` and `73.856 us` for joined exact FP16. Joined FP13 is + also slightly slower than three-stream FP13. Spending another 20.5 MiB per + C4 layer for a slower current-source route is rejected. +- The previously excluded shared-expert gate/up role was re-audited without + the unsafe fused-SiLU epilogue. The exact production semantics are full + `K4096/N1024` gate/up followed by `SiluAndMulWithClamp(10)`. Non-fused QPN8 + split-32/two-chain plus the identical clamp moves `62.299` to `11.582 us` + (`5.38x`) on real weights. Two low-scale patterns are bitwise after clamp; + the remaining patterns retain at least `99.609%` exact FP16 elements with + maximum relative L2 `4.63e-6`. The exact tensor-role route retains the + external production activation and does not inspect model or checkpoint + identity, but the operator bound alone is not sufficient for promotion. + `VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE=1` is an experimental opt-in. +- The matched shared-gate endpoint rejects that opt-in on the quality gate. + Control records a median `75.166 token/s` (`13.304 ms/token`) and GSM8K-64 + `62/64`; the candidate records `75.641 token/s` (`13.220 ms/token`) but + scores `61/64`. Only item 37 changes correctness (`2 -> 0`); 35/64 texts + and 62/64 extracted numeric predictions are identical, with no candidate + win. The A/B changes only the shared-expert gate/up projection while keeping + the external clamp-SwiGLU identical, isolating the greedy-trajectory change + to the QPN8 projection/reduction path. The `0.084 ms/token` speed reduction + is therefore not admitted, and the route now defaults off while an exact + TurboMind-prescaled alternative is screened. +- The replacement keeps the gated-interleaved layout, layout restore, and + external clamp-SwiGLU, and its FP16 block-scale exponent shift is reversible. + It does not keep the control accumulation path: C++ dispatch selects the + dedicated `Config_E4M3_Prescaled` `8x128x64`, split-K-7, swizzle-1 kernel, + changing partial boundaries and reduction order versus the ordinary + TurboMind control. The earlier same-accumulation description is withdrawn. + On the real layer-0 TP4-rank-0 tensor, the initial 64 dynamic patterns happen + to be bitwise before and after activation and both CUDA Graphs replay stably; + same-GPU A/B/B/A timing moves `62.298` to `20.562 us/layer` (`3.03x`). This + screen is insufficient for inputs near FP16 rounding boundaries. Evidence + is under + `/data/models/v100-dsv4-0731-pp2tp4-shared-gate-prescaled-exact-screen-20260826-r1/`. + The isolated layer-0 result does not generalize to the model endpoint. In a + source/topology/request-matched full pair, the control scores GSM8K `64/64`, + HumanEval `29/32`, LongBench `44.740`, and `73.539 token/s`; the candidate + scores `62/64`, `28/32`, and has LongBench row regressions while slowing to + `73.302 token/s`. GSM items 37 and 54 regress with no wins, HumanEval/16 + regresses, and all three performance streams change. The route is rejected + and its runtime integration is removed rather than retained as an opt-in. + Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-private-quality-safe-prescaled-shared-gate-{control,candidate}-fullmodel-20260826-r2/`. +- Sparse-MLA long-context decay is now directly quantified. With the current + exact `BLOCK_K=16`, one C128 layer measures `12.43 us` at 1K context, + `18.37 us` at 64K, `26.98 us` at 128K, `42.40 us` at 256K, + `85.63 us` at 512K and `168.74 us` at 1M. `BLOCK_K=32/64` is slower at every + tested context and changes roughly 36--44% of FP16 elements; `BLOCK_K=8` + is not a valid Triton dot shape. Dynamic larger-K selection is rejected. + Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-sparse-mla-blockk-screen-20260826-r1/`. +- Grouping two independent 64-D value tiles per CTA is bitwise, but its gain is + only `1.25%` at C4, `0.78%` at 64K and `1.21%` at 1M, while SWA and 1K + regress; four/eight-tile waves regress further. The D-wave source candidate + is therefore not retained. Evidence is under + `/data/models/v100-dsv4-0731-pp2tp4-sparse-mla-dwave-screen-20260826-r1/`. +- A primary-source long-context review confirms that generic Flash-Decoding + is not a missing route: the official FlashAttention inference path splits + KV sequence loading across CTAs and combines online-softmax partials, while + FlashInfer sizes KV chunks against available SM work. The current V4 SM70 + decode kernel already implements the same structure with direct FP8 paged + loads, 16-token partial max/sum/accumulators and a separate reduction. + Local K32/K64 and D-wave results therefore reject a blind copy of either + generic scheduler. The next long-context trace must separate QK tiles, + score reduction, value tiles and final reduction. Only if the merge or SM + tail is dominant should a LeanAttention-style hierarchical stream-K merge + be screened above a long-context threshold; TP4-local decode-context + partitioning over NVLink is the alternative when KV bandwidth dominates. + Primary references are the + [FlashAttention inference implementation](https://github.com/Dao-AILab/flash-attention/blob/main/README.md), + [FlashInfer scheduler](https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/attention/scheduler.cuh), + [FlashInfer paper](https://arxiv.org/abs/2501.01005), and + [LeanAttention paper](https://arxiv.org/abs/2405.10480). +- The next bounded MoE screen replaces the historical nondeterministic FP16 + atomic weighted epilogue with one CTA per W2 `N=32` tile. It rounds each + route result to FP16, then performs the same fixed route-order FP32 FMA as + the standalone reducer. The SM70 AOT build passes; the exclusive-V100 + bitwise/performance screen is queued only after the full-model A/B pair. + No production source or default changes are admitted from this build-only + result. +- Shared-machine GPU ownership remains fail-closed: benchmarks wait on pidfds, + claim only genuinely idle cards, never terminate unrelated processes, and + release claims immediately after bounded tests. Development and any Draft + PR for this continuation use the private remote and the isolated branch + `agent/private-v100-dsv4-pp2tp4-followup-20260826`; public upstream is + fetch-only for this work. + +## 2026-08-27 NVFP4 DFlash2 output-quality audit + +- This audit starts from private main + `d63e9490f65f9e01f6649053c1ab72922034b931` in isolated worktree + `v100-dflash2-quality-audit-20260826-124851`. Source changes are in private + Draft PR #18; no public branch is used. Remote deployment and packaging are + explicitly out of scope. +- The packaged launch forced selector QPN8 candidate-order tie resolution even + though source defaults retained dense-vocabulary order. Real-hidden evidence + shows a 1.49% top-20 set change and rare top-1 changes at equal-valued + cutoffs. Candidate order now requires a separate benchmark-only opt-in; + stale `DENSE_ORDER=0` launch files are ignored and log `dense_order=True`. + The five-arm three-seed screen keeps D1 and D2 at `14/24`; D2 improves mean + steady decode by 3.88% and changes request acceptance by `-0.0195`. The + unsafe D3 arm scores `13/24` and is not promoted. +- The local fork predated upstream structured-output fixes. The minimal + backport uses the scheduler's exact accepted tokens, trims reasoning before + grammar advance, constrains bonus/post-marker rows, avoids expected FSM + errors for drafts produced before their masks, and stops/resets XGrammar at + termination. It does not change Eagle, MTP, DDTree, or their routing. +- The practical temperature-1.0 coding gate uses NVFP4 target, official BF16 + DFlash2 draft, TP4 V100, target E5M2 KV, draft FP16 KV, 256K max context, + 4096 chunking, prefix cache, Mamba align, probabilistic block-8, and full + Graphs. Three seed bases produce: + - MBPP EvalPlus Base `89/93`, Plus `80/93`, natural stop `95/96`; + - HumanEval Base `94/96`, Plus `92/96`, natural stop `91/96`; + - LiveCodeBench `33/48`, exactly `11/16` for each seed, natural stop `32/48`. + The first-seed MBPP+HumanEval aggregate exactly matches historical + no-DFlash at Base `62/63`, Plus `59/63`. LiveCodeBench also matches both + historical no-DFlash and old DFlash at `11/16`, exchanging one passing task + rather than changing the total score. +- LiveCodeBench's aggregate output rate is `170.283 token/s`, mean per-request + steady decode is `196.541 token/s`, pooled acceptance is `3.277`, and mean + per-request acceptance is `3.694`. All failures in the first two seeds are + 16K length cases, while all naturally stopped cases pass. Excessive thinking + at temperature 1.0, not verifier corruption, is the remaining coding-product + risk. +- The same weak middle seed at the official precise-coding temperature 0.6 + improves MBPP Base/Plus from `27/24` to `29/27` out of 31, keeps HumanEval + Base/Plus at `31/31` out of 32, and keeps LiveCodeBench at `11/16`. Across all + 80 requests, request-mean acceptance moves `4.273 -> 4.514` and mean steady + decode `233.187 -> 244.520 token/s`, but natural stops fall from `72/80` to + `70/80`. Temperature 0.6 is therefore an optional precise-coding profile, + not a global API-default replacement for temperature 1.0. +- The current 256K/4096/prefix/Mamba-align Graph PPL pair scores 16,376 fixed + WikiText tokens. Target-only/DFlash2 weighted PPL is + `5.4993116/5.4993622`, a `+0.0000506` (`+0.00092%`) change; maximum per-segment + difference is `0.0062143`. This passes the fixed 0.01 segment-PPL bound and + provides distribution-level evidence against DFlash2-induced target-model + degradation. +- The actual 256K Graph API gate enables prefix cache, Mamba align, `qwen3` + reasoning and `qwen3_coder` tools. B1 and B4 each pass 12/12 across + `json_object`, strict schema, and required tool calls. Alternating long + ALPHA/BETA prefixes pass 5/5 with exact sentinels and checksums; metrics show + 59,396 queried and 32,960 cached tokens. The error-signature scan is empty. +- Evidence is rooted at + `/data/minimax-h3/task-cache/v100-dflash2-quality-audit-20260826/`. + Detailed rationale and per-arm results are in + `docs/design/sm70_dflash2_quality_audit.md`. Do not cite acceptance length as + intelligence: it explains throughput, while official executable scores and + natural termination are the quality gates. diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index bda540b88c..23e6af365e 100644 --- a/docs/features/speculative_decoding/README.md +++ b/docs/features/speculative_decoding/README.md @@ -119,30 +119,45 @@ vllm serve \ #### DFlash2 with n-gram assistance -Set `ngram_assist` on a DFlash2 configuration to try prompt lookup before the -neural DFlash2 query and selector. When every active request has a full-width -lookup hit, the batch skips that neural proposal work. Mixed-hit and miss -batches keep the unchanged DFlash2 proposals. The target model still verifies -every proposed token, including probabilistic sampling, so this does not -change the target sampling distribution. +Set `ngram_assist` on a DFlash2 configuration to combine prompt lookup with +the neural drafter. Two layouts are supported: + +- At the checkpoint-native width (seven drafts for the official block-8 + checkpoint), a full-batch lookup hit may skip the neural query and selector. +- With a wider configured verifier, DFlash2 still produces only its seven + trained drafts. GPU lookup can replace that prefix and fill the remaining + positions from the request's own token history. The scheduler normally + verifies q8 and switches to q16 only after two consecutive strong copy + steps. + +The target model verifies every scheduled proposal. In probabilistic mode, +each lookup position is represented by a point-mass draft distribution, so +rejection sampling preserves the target distribution. ```bash vllm serve \ --speculative-config '{ "method": "dflash", "model": "your-org/your-dflash2-selector-model", - "num_speculative_tokens": 7, + "num_speculative_tokens": 15, "draft_sample_method": "probabilistic", "ngram_assist": true, - "prompt_lookup_min": 5, - "prompt_lookup_max": 5 + "prompt_lookup_min": 6, + "prompt_lookup_max": 12 }' ``` This hybrid route currently requires `method=dflash` and a DFlash2 checkpoint. It is intentionally not enabled for DFlash1, Eagle, MTP, or -`dflash_ddtree`. Structured-output requests bypass prompt lookup and continue -through DFlash2 until grammar-aware proposal masking is available. +`dflash_ddtree`. Structured-output batches bypass lookup and retain the +checkpoint-native q8 DFlash2 path until grammar-aware proposal masking is +available. + +Adaptive q8/q16 verification requires the synchronous scheduler because the +asynchronous scheduler fixes the proposal width at startup. vLLM disables +asynchronous scheduling automatically for this layout; explicitly enabling it +raises a configuration error. Set `VLLM_DFLASH2_LOOKUP_ADAPTIVE=0` only when a +fixed q16 verifier is deliberately preferred. #### Suffix decoding diff --git a/flash-attention-v100/flash_attn_v100/flash_attn_interface.py b/flash-attention-v100/flash_attn_v100/flash_attn_interface.py index c7042bbdb3..a0ae0c4db6 100644 --- a/flash-attention-v100/flash_attn_v100/flash_attn_interface.py +++ b/flash-attention-v100/flash_attn_v100/flash_attn_interface.py @@ -1036,11 +1036,12 @@ def flash_attn_grouped_verify_paged( v_scale: float = 1.0, one_pass: bool = False, ) -> torch.Tensor: - """Exact grouped q8/H6/D256 DFlash2 verifier prototype for SM70. + """Exact grouped q8/q16 H6/D256 DFlash2 verifier for SM70. The native entry keeps all causal verifier rows together and reuses each - paged-KV scan across four GQA heads. Workspaces are stream-local and fixed - size so a warmed entry is CUDA-graph safe. Unsupported shapes fail closed. + paged-KV scan across a packed GQA group. q8 uses one six-head group and q16 + uses two three-head groups, retaining 48 rows per CTA and the same workspace + byte count. Workspaces are stream-local and CUDA-graph safe. """ if softmax_scale is None: softmax_scale = q.shape[-1] ** -0.5 @@ -1049,6 +1050,12 @@ def flash_attn_grouped_verify_paged( seq_lens = maybe_contiguous(seq_lens) out = maybe_contiguous(out) workspace = _get_grouped_verify_workspace(q) + if q.shape[0] > 8: + partial_out = workspace.partial_out.view(40, 16, 6, 256) + partial_lse = workspace.partial_lse.view(40, 16, 6) + else: + partial_out = workspace.partial_out + partial_lse = workspace.partial_lse return flash_attn_v100_cuda.grouped_verify_paged_fwd( q, k_cache, @@ -1056,8 +1063,8 @@ def flash_attn_grouped_verify_paged( out, block_table, seq_lens, - workspace.partial_out, - workspace.partial_lse, + partial_out, + partial_lse, float(softmax_scale), kv_cache_dtype, float(k_scale), diff --git a/flash-attention-v100/include/fused_mma.h b/flash-attention-v100/include/fused_mma.h index 2d7c89f883..407289c504 100644 --- a/flash-attention-v100/include/fused_mma.h +++ b/flash-attention-v100/include/fused_mma.h @@ -2,7 +2,7 @@ #define FUSED_MMA_H #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ != 700) -#error "Volta WMMA: This header is for sm_70 ONLY! Compile with -arch=sm_70" + #error "Volta WMMA: This header is for sm_70 ONLY! Compile with -arch=sm_70" #endif #include @@ -15,385 +15,409 @@ struct matrix_a {}; struct matrix_b {}; struct accumulator {}; -enum layout_t { - mem_row_major, - mem_col_major -}; +enum layout_t { mem_row_major, mem_col_major }; template struct fragment; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; -template <> struct fragment { uint32_t x[8]; static constexpr int num_elements = 16; }; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; +template <> +struct fragment { + uint32_t x[8]; + static constexpr int num_elements = 16; +}; -template <> struct fragment { float x[8]; static constexpr int num_elements = 8; }; -template <> struct fragment { float x[8]; static constexpr int num_elements = 8; }; -template <> struct fragment { float x[8]; static constexpr int num_elements = 8; }; +template <> +struct fragment { + float x[8]; + static constexpr int num_elements = 8; +}; +template <> +struct fragment { + float x[8]; + static constexpr int num_elements = 8; +}; +template <> +struct fragment { + float x[8]; + static constexpr int num_elements = 8; +}; __device__ __forceinline__ unsigned get_lane_id() { - unsigned lane_id; - asm volatile ("mov.u32 %0, %%laneid;" : "=r"(lane_id)); - return lane_id; + unsigned lane_id; + asm volatile("mov.u32 %0, %%laneid;" : "=r"(lane_id)); + return lane_id; } template -__device__ __forceinline__ void fill_fragment(fragment& frag, float value) { - #pragma unroll - for (int i = 0; i < 8; ++i) frag.x[i] = value; +__device__ __forceinline__ void fill_fragment( + fragment& frag, float value) { +#pragma unroll + for (int i = 0; i < 8; ++i) frag.x[i] = value; } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { - asm volatile( - "wmma.load.a.sync.aligned.row.m16n16k16.f16 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.a.sync.aligned.row.m16n16k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { - asm volatile( - "wmma.load.a.sync.aligned.col.m16n16k16.f16 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.a.sync.aligned.col.m16n16k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { - asm volatile( - "wmma.load.b.sync.aligned.row.m16n16k16.f16 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.b.sync.aligned.row.m16n16k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.b.sync.aligned.col.m16n16k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); +} + +__device__ __forceinline__ void load_matrix_sync( + fragment& frag, const float* smem_ptr, + unsigned ldm, layout_t layout) { + if (layout == mem_row_major) { + asm volatile( + "wmma.load.c.sync.aligned.row.m16n16k16.f32 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), + "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); + } else { asm volatile( - "wmma.load.b.sync.aligned.col.m16n16k16.f16 " + "wmma.load.c.sync.aligned.col.m16n16k16.f32 " "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), + "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + : "memory"); + } } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const float* smem_ptr, unsigned ldm, layout_t layout) { - if (layout == mem_row_major) { - asm volatile( - "wmma.load.c.sync.aligned.row.m16n16k16.f32 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), - "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); - } else { - asm volatile( - "wmma.load.c.sync.aligned.col.m16n16k16.f32 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), - "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); - } + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.a.sync.aligned.row.m32n8k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { - asm volatile( - "wmma.load.a.sync.aligned.row.m32n8k16.f16 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.a.sync.aligned.col.m32n8k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { - asm volatile( - "wmma.load.a.sync.aligned.col.m32n8k16.f16 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.b.sync.aligned.row.m32n8k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { - asm volatile( - "wmma.load.b.sync.aligned.row.m32n8k16.f16 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.b.sync.aligned.col.m32n8k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { + fragment& frag, const float* smem_ptr, + unsigned ldm, layout_t layout) { + if (layout == mem_row_major) { asm volatile( - "wmma.load.b.sync.aligned.col.m32n8k16.f16 " + "wmma.load.c.sync.aligned.row.m32n8k16.f32 " "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), + "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + : "memory"); + } else { + asm volatile( + "wmma.load.c.sync.aligned.col.m32n8k16.f32 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), + "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); + } } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const float* smem_ptr, unsigned ldm, layout_t layout) { - if (layout == mem_row_major) { - asm volatile( - "wmma.load.c.sync.aligned.row.m32n8k16.f32 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), - "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); - } else { - asm volatile( - "wmma.load.c.sync.aligned.col.m32n8k16.f32 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), - "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); - } + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.a.sync.aligned.row.m8n32k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { - asm volatile( - "wmma.load.a.sync.aligned.row.m8n32k16.f16 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.a.sync.aligned.col.m8n32k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { - asm volatile( - "wmma.load.a.sync.aligned.col.m8n32k16.f16 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.b.sync.aligned.row.m8n32k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { - asm volatile( - "wmma.load.b.sync.aligned.row.m8n32k16.f16 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); + fragment& frag, const half* smem_ptr, + unsigned ldm) { + asm volatile( + "wmma.load.b.sync.aligned.col.m8n32k16.f16 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), + "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); } __device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const half* smem_ptr, unsigned ldm) { + fragment& frag, const float* smem_ptr, + unsigned ldm, layout_t layout) { + if (layout == mem_row_major) { asm volatile( - "wmma.load.b.sync.aligned.col.m8n32k16.f16 " + "wmma.load.c.sync.aligned.row.m8n32k16.f32 " "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=r"(frag.x[0]), "=r"(frag.x[1]), "=r"(frag.x[2]), "=r"(frag.x[3]), - "=r"(frag.x[4]), "=r"(frag.x[5]), "=r"(frag.x[6]), "=r"(frag.x[7]) + : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), + "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) : "l"(smem_ptr), "r"(ldm) - : "memory" - ); -} - -__device__ __forceinline__ void load_matrix_sync( - fragment& frag, - const float* smem_ptr, unsigned ldm, layout_t layout) { - if (layout == mem_row_major) { - asm volatile( - "wmma.load.c.sync.aligned.row.m8n32k16.f32 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), - "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); - } else { - asm volatile( - "wmma.load.c.sync.aligned.col.m8n32k16.f32 " - "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" - : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), - "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) - : "l"(smem_ptr), "r"(ldm) - : "memory" - ); - } + : "memory"); + } else { + asm volatile( + "wmma.load.c.sync.aligned.col.m8n32k16.f32 " + "{%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=f"(frag.x[0]), "=f"(frag.x[1]), "=f"(frag.x[2]), "=f"(frag.x[3]), + "=f"(frag.x[4]), "=f"(frag.x[5]), "=f"(frag.x[6]), "=f"(frag.x[7]) + : "l"(smem_ptr), "r"(ldm) + : "memory"); + } } __device__ __forceinline__ void store_matrix_sync( - float* smem_ptr, - const fragment& frag, + float* smem_ptr, const fragment& frag, unsigned ldm, layout_t layout) { - if (layout == mem_row_major) { - asm volatile( - "wmma.store.d.sync.aligned.row.m16n16k16.f32 " - "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" - : - : "l"(smem_ptr), - "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), "f"(frag.x[3]), - "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), "f"(frag.x[7]), - "r"(ldm) - : "memory" - ); - } else { - asm volatile( - "wmma.store.d.sync.aligned.col.m16n16k16.f32 " - "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" - : - : "l"(smem_ptr), - "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), "f"(frag.x[3]), - "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), "f"(frag.x[7]), - "r"(ldm) - : "memory" - ); - } + if (layout == mem_row_major) { + asm volatile( + "wmma.store.d.sync.aligned.row.m16n16k16.f32 " + "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" + : + : "l"(smem_ptr), "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), + "f"(frag.x[3]), "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), + "f"(frag.x[7]), "r"(ldm) + : "memory"); + } else { + asm volatile( + "wmma.store.d.sync.aligned.col.m16n16k16.f32 " + "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" + : + : "l"(smem_ptr), "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), + "f"(frag.x[3]), "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), + "f"(frag.x[7]), "r"(ldm) + : "memory"); + } } __device__ __forceinline__ void store_matrix_sync( - float* smem_ptr, - const fragment& frag, + float* smem_ptr, const fragment& frag, unsigned ldm, layout_t layout) { - if (layout == mem_row_major) { - asm volatile( - "wmma.store.d.sync.aligned.row.m32n8k16.f32 " - "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" - : - : "l"(smem_ptr), - "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), "f"(frag.x[3]), - "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), "f"(frag.x[7]), - "r"(ldm) - : "memory" - ); - } else { - asm volatile( - "wmma.store.d.sync.aligned.col.m32n8k16.f32 " - "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" - : - : "l"(smem_ptr), - "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), "f"(frag.x[3]), - "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), "f"(frag.x[7]), - "r"(ldm) - : "memory" - ); - } + if (layout == mem_row_major) { + asm volatile( + "wmma.store.d.sync.aligned.row.m32n8k16.f32 " + "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" + : + : "l"(smem_ptr), "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), + "f"(frag.x[3]), "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), + "f"(frag.x[7]), "r"(ldm) + : "memory"); + } else { + asm volatile( + "wmma.store.d.sync.aligned.col.m32n8k16.f32 " + "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" + : + : "l"(smem_ptr), "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), + "f"(frag.x[3]), "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), + "f"(frag.x[7]), "r"(ldm) + : "memory"); + } } __device__ __forceinline__ void store_matrix_sync( - float* smem_ptr, - const fragment& frag, + float* smem_ptr, const fragment& frag, unsigned ldm, layout_t layout) { - if (layout == mem_row_major) { - asm volatile( - "wmma.store.d.sync.aligned.row.m8n32k16.f32 " - "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" - : - : "l"(smem_ptr), - "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), "f"(frag.x[3]), - "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), "f"(frag.x[7]), - "r"(ldm) - : "memory" - ); - } else { - asm volatile( - "wmma.store.d.sync.aligned.col.m8n32k16.f32 " - "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" - : - : "l"(smem_ptr), - "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), "f"(frag.x[3]), - "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), "f"(frag.x[7]), - "r"(ldm) - : "memory" - ); - } + if (layout == mem_row_major) { + asm volatile( + "wmma.store.d.sync.aligned.row.m8n32k16.f32 " + "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" + : + : "l"(smem_ptr), "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), + "f"(frag.x[3]), "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), + "f"(frag.x[7]), "r"(ldm) + : "memory"); + } else { + asm volatile( + "wmma.store.d.sync.aligned.col.m8n32k16.f32 " + "[%0], {%1,%2,%3,%4,%5,%6,%7,%8}, %9;" + : + : "l"(smem_ptr), "f"(frag.x[0]), "f"(frag.x[1]), "f"(frag.x[2]), + "f"(frag.x[3]), "f"(frag.x[4]), "f"(frag.x[5]), "f"(frag.x[6]), + "f"(frag.x[7]), "r"(ldm) + : "memory"); + } } -#define VOLTA_WMMA_MMA_F32(M, N, K, ALAY, BLAY) \ -__device__ __forceinline__ void mma_sync( \ - fragment& d, \ - const fragment& a, \ - const fragment& b, \ - const fragment& c) { \ - asm volatile( \ - "wmma.mma.sync.aligned." #ALAY "." #BLAY ".m" #M "n" #N "k" #K ".f32.f32 " \ - "{%0,%1,%2,%3,%4,%5,%6,%7}, " \ - "{%8,%9,%10,%11,%12,%13,%14,%15}, " \ - "{%16,%17,%18,%19,%20,%21,%22,%23}, " \ - "{%24,%25,%26,%27,%28,%29,%30,%31};" \ - : "=f"(d.x[0]), "=f"(d.x[1]), "=f"(d.x[2]), "=f"(d.x[3]), \ - "=f"(d.x[4]), "=f"(d.x[5]), "=f"(d.x[6]), "=f"(d.x[7]) \ - : "r"(a.x[0]), "r"(a.x[1]), "r"(a.x[2]), "r"(a.x[3]), \ - "r"(a.x[4]), "r"(a.x[5]), "r"(a.x[6]), "r"(a.x[7]), \ - "r"(b.x[0]), "r"(b.x[1]), "r"(b.x[2]), "r"(b.x[3]), \ - "r"(b.x[4]), "r"(b.x[5]), "r"(b.x[6]), "r"(b.x[7]), \ - "f"(c.x[0]), "f"(c.x[1]), "f"(c.x[2]), "f"(c.x[3]), \ - "f"(c.x[4]), "f"(c.x[5]), "f"(c.x[6]), "f"(c.x[7]) \ - ); \ -} +#define VOLTA_WMMA_MMA_F32(M, N, K, ALAY, BLAY) \ + __device__ __forceinline__ void mma_sync( \ + fragment& d, \ + const fragment& a, \ + const fragment& b, \ + const fragment& c) { \ + asm volatile( \ + "wmma.mma.sync.aligned." #ALAY "." #BLAY ".m" #M "n" #N "k" #K \ + ".f32.f32 " \ + "{%0,%1,%2,%3,%4,%5,%6,%7}, " \ + "{%8,%9,%10,%11,%12,%13,%14,%15}, " \ + "{%16,%17,%18,%19,%20,%21,%22,%23}, " \ + "{%24,%25,%26,%27,%28,%29,%30,%31};" \ + : "=f"(d.x[0]), "=f"(d.x[1]), "=f"(d.x[2]), "=f"(d.x[3]), \ + "=f"(d.x[4]), "=f"(d.x[5]), "=f"(d.x[6]), "=f"(d.x[7]) \ + : "r"(a.x[0]), "r"(a.x[1]), "r"(a.x[2]), "r"(a.x[3]), "r"(a.x[4]), \ + "r"(a.x[5]), "r"(a.x[6]), "r"(a.x[7]), "r"(b.x[0]), "r"(b.x[1]), \ + "r"(b.x[2]), "r"(b.x[3]), "r"(b.x[4]), "r"(b.x[5]), "r"(b.x[6]), \ + "r"(b.x[7]), "f"(c.x[0]), "f"(c.x[1]), "f"(c.x[2]), "f"(c.x[3]), \ + "f"(c.x[4]), "f"(c.x[5]), "f"(c.x[6]), "f"(c.x[7])); \ + } VOLTA_WMMA_MMA_F32(16, 16, 16, row, col) VOLTA_WMMA_MMA_F32(16, 16, 16, row, row) @@ -412,6 +436,6 @@ VOLTA_WMMA_MMA_F32(8, 32, 16, col, row) #undef VOLTA_WMMA_MMA_F32 -} +} // namespace volta #endif diff --git a/flash-attention-v100/kernel/contiguous_to_paged.cu b/flash-attention-v100/kernel/contiguous_to_paged.cu index 1c29b0678d..a6561485e7 100644 --- a/flash-attention-v100/kernel/contiguous_to_paged.cu +++ b/flash-attention-v100/kernel/contiguous_to_paged.cu @@ -4,89 +4,72 @@ #include __global__ void contiguous_to_paged_kernel( - const __half* __restrict__ key, - const __half* __restrict__ value, - const int64_t* __restrict__ slot_mapping, - __half* __restrict__ key_cache, - __half* __restrict__ value_cache, - const int64_t key_stride, - const int64_t value_stride, - const int64_t block_stride, - const int64_t page_stride, - const int64_t head_stride, - const int block_size, - const int num_heads, - const int head_dim -) { + const __half* __restrict__ key, const __half* __restrict__ value, + const int64_t* __restrict__ slot_mapping, __half* __restrict__ key_cache, + __half* __restrict__ value_cache, const int64_t key_stride, + const int64_t value_stride, const int64_t block_stride, + const int64_t page_stride, const int64_t head_stride, const int block_size, + const int num_heads, const int head_dim) { + const int token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[token_idx]; - const int token_idx = blockIdx.x; - const int64_t slot_idx = slot_mapping[token_idx]; + if (slot_idx < 0) { + return; + } - if (slot_idx < 0) { - return; - } - - const int64_t block_idx = slot_idx / block_size; - const int64_t block_offset = slot_idx % block_size; + const int64_t block_idx = slot_idx / block_size; + const int64_t block_offset = slot_idx % block_size; - const int tid = threadIdx.x; - const int num_threads = blockDim.x; + const int tid = threadIdx.x; + const int num_threads = blockDim.x; - const __half* __restrict__ key_src = key + token_idx * key_stride; - const __half* __restrict__ value_src = value + token_idx * value_stride; + const __half* __restrict__ key_src = key + token_idx * key_stride; + const __half* __restrict__ value_src = value + token_idx * value_stride; - __half* __restrict__ key_dst = key_cache + block_idx * block_stride + block_offset * page_stride; - __half* __restrict__ value_dst = value_cache + block_idx * block_stride + block_offset * page_stride; + __half* __restrict__ key_dst = + key_cache + block_idx * block_stride + block_offset * page_stride; + __half* __restrict__ value_dst = + value_cache + block_idx * block_stride + block_offset * page_stride; - for (int head_idx = 0; head_idx < num_heads; head_idx++) { - const __half* key_head_src = key_src + head_idx * head_dim; - const __half* value_head_src = value_src + head_idx * head_dim; - __half* key_head_dst = key_dst + head_idx * head_stride; - __half* value_head_dst = value_dst + head_idx * head_stride; + for (int head_idx = 0; head_idx < num_heads; head_idx++) { + const __half* key_head_src = key_src + head_idx * head_dim; + const __half* value_head_src = value_src + head_idx * head_dim; + __half* key_head_dst = key_dst + head_idx * head_stride; + __half* value_head_dst = value_dst + head_idx * head_stride; - for (int elem_idx = tid; elem_idx < head_dim; elem_idx += num_threads) { - key_head_dst[elem_idx] = key_head_src[elem_idx]; - value_head_dst[elem_idx] = value_head_src[elem_idx]; - } + for (int elem_idx = tid; elem_idx < head_dim; elem_idx += num_threads) { + key_head_dst[elem_idx] = key_head_src[elem_idx]; + value_head_dst[elem_idx] = value_head_src[elem_idx]; } + } } -void contiguous_to_paged( - torch::Tensor key, - torch::Tensor value, - torch::Tensor slot_mapping, - torch::Tensor key_cache, - torch::Tensor value_cache -) { - const int num_tokens = slot_mapping.size(0); - const int num_heads = key.size(1); - const int head_dim = key.size(2); - const int block_size = key_cache.size(1); +void contiguous_to_paged(torch::Tensor key, torch::Tensor value, + torch::Tensor slot_mapping, torch::Tensor key_cache, + torch::Tensor value_cache) { + const int num_tokens = slot_mapping.size(0); + const int num_heads = key.size(1); + const int head_dim = key.size(2); + const int block_size = key_cache.size(1); - const int64_t key_stride = key.stride(0); - const int64_t value_stride = value.stride(0); - const int64_t block_stride = key_cache.stride(0); - const int64_t page_stride = key_cache.stride(1); - const int64_t head_stride = key_cache.stride(2); + const int64_t key_stride = key.stride(0); + const int64_t value_stride = value.stride(0); + const int64_t block_stride = key_cache.stride(0); + const int64_t page_stride = key_cache.stride(1); + const int64_t head_stride = key_cache.stride(2); - const int threads = 256; - contiguous_to_paged_kernel<<>>( - reinterpret_cast(key.data_ptr()), - reinterpret_cast(value.data_ptr()), - slot_mapping.data_ptr(), - reinterpret_cast<__half*>(key_cache.data_ptr()), - reinterpret_cast<__half*>(value_cache.data_ptr()), - key_stride, - value_stride, - block_stride, - page_stride, - head_stride, - block_size, - num_heads, - head_dim - ); + const int threads = 256; + contiguous_to_paged_kernel<<>>( + reinterpret_cast(key.data_ptr()), + reinterpret_cast(value.data_ptr()), + slot_mapping.data_ptr(), + reinterpret_cast<__half*>(key_cache.data_ptr()), + reinterpret_cast<__half*>(value_cache.data_ptr()), key_stride, + value_stride, block_stride, page_stride, head_stride, block_size, + num_heads, head_dim); } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("contiguous_to_paged", &contiguous_to_paged, "Contiguous K/V to Paged KV Cache"); + m.def("contiguous_to_paged", &contiguous_to_paged, + "Contiguous K/V to Paged KV Cache"); } diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index b6417092d4..6d92630656 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -211,6 +211,16 @@ bool xqa_e5m2_batch_wide_load_enabled() { return value == nullptr || value[0] != '0'; } +bool dflash2_grouped_fixed_interleaved_enabled() { + const char* value = std::getenv("VLLM_FLASH_V100_DFLASH2_FIXED_INTERLEAVED"); + return value == nullptr || value[0] != '0'; +} + +bool dflash2_grouped_stage_page_ids_enabled() { + const char* value = std::getenv("VLLM_FLASH_V100_DFLASH2_STAGE_PAGE_IDS"); + return value == nullptr || value[0] != '0'; +} + int xqa_e5m2_p1024_begin() { const char* value = std::getenv("VLLM_FLASH_V100_XQA_E5M2_P1024_BEGIN"); return value == nullptr ? 61633 : std::max(1, std::atoi(value)); @@ -737,7 +747,8 @@ __device__ __forceinline__ uint4 load_xqa_tc_kv_vector( BLOCK_SIZE == 1648 || BLOCK_SIZE == 3296, "Unsupported paged-KV block-size specialization"); static_assert(!CONTIGUOUS_HKV1_LAYOUT || BLOCK_SIZE == 16 || - BLOCK_SIZE == 800 || BLOCK_SIZE == 1568, + BLOCK_SIZE == 800 || BLOCK_SIZE == 1568 || + BLOCK_SIZE == 1648 || BLOCK_SIZE == 3296, "The fixed-stride Hkv=1 layout requires a specialized page"); int logical_block; int block_offset; @@ -856,7 +867,7 @@ __device__ __forceinline__ void load_xqa_tc_kv_panel( if constexpr (CONTIGUOUS_HKV1_LAYOUT) { constexpr int64_t kHeadDim = 256; constexpr int64_t kPhysicalBlockStride = - BLOCK_SIZE == 16 ? 16 * kHeadDim : 2 * 800 * kHeadDim; + BLOCK_SIZE == 16 ? 16 * kHeadDim : 2 * BLOCK_SIZE * kHeadDim; physical_offset = static_cast(physical_block) * kPhysicalBlockStride + static_cast(block_offset) * kHeadDim + panel_offset; @@ -1766,21 +1777,28 @@ __global__ void __launch_bounds__(NUM_THREADS, MIN_BLOCKS_PER_SM) // It remains a numerical oracle for the production one-pass register-rescale // variant, which loads K once and preserves the same online-softmax/PV result. // -// The verifier shape has one KV head and six query heads. Pack all six heads -// for all eight verifier tokens into one 48-row CTA so each context split -// loads and decodes K/V once. This is exact GQA packing: only the CTA mapping -// changes; QK, online softmax, PV, and split-state reduction stay unchanged. -constexpr int kGroupedVerifyMaxQ = 8; +// The verifier shape has one KV head and six query heads. Keep each CTA at 48 +// rows for both supported widths: q8 packs six heads in one CTA, while q16 +// packs three heads in each of two CTAs. The q16 route halves the number of +// context splits, preserving eighty CTAs and the q8 workspace byte count. +constexpr int kGroupedVerifyQ8MaxQ = 8; +constexpr int kGroupedVerifyQ16MaxQ = 16; +constexpr int kGroupedVerifyMaxSupportedQ = kGroupedVerifyQ16MaxQ; constexpr int kGroupedVerifyHeads = 6; constexpr int kGroupedVerifyHeadDim = 256; -constexpr int kGroupedVerifyHeadsPerCta = 6; -constexpr int kGroupedVerifyRows = - kGroupedVerifyMaxQ * kGroupedVerifyHeadsPerCta; +constexpr int kGroupedVerifyRows = 48; constexpr int kGroupedVerifyBlockN = 32; +constexpr int kGroupedVerifyQStride = 264; +constexpr int kGroupedVerifyKVStride = 264; +constexpr int kGroupedVerifyScoreStride = 32; +constexpr int kGroupedVerifyProbStride = 40; +constexpr int kGroupedVerifyPageIdsCapacity = 16; // One packed head group times eighty context splits maps one 512-thread CTA to // each of V100's eighty SMs. Short sequences still reduce active_splits using // kGroupedVerifyMinTokensPerSplit inside the captured graph. -constexpr int kGroupedVerifySplits = 80; +constexpr int kGroupedVerifyQ8Splits = 80; +constexpr int kGroupedVerifyWorkspaceRows = + kGroupedVerifyQ8MaxQ * kGroupedVerifyQ8Splits; // A packed CTA replaces two 256-thread head-group CTAs and can occupy only one // SM. Use 64-token chunks to expose enough short-context parallelism without // paying one split reduction per N32 tile; longer contexts saturate at the @@ -1798,19 +1816,36 @@ constexpr int kGroupedVerifyOutputTiles = constexpr int kGroupedVerifyOutputTilesPerWarp = kGroupedVerifyOutputTiles / kGroupedVerifyWarps; +template +struct GroupedVerifyTraits { + static_assert(MAX_QUERY_TOKENS == kGroupedVerifyQ8MaxQ || + MAX_QUERY_TOKENS == kGroupedVerifyQ16MaxQ, + "grouped verifier supports q8 and q16 workspaces only"); + static constexpr int kHeadsPerCta = kGroupedVerifyRows / MAX_QUERY_TOKENS; + static constexpr int kHeadGroups = kGroupedVerifyHeads / kHeadsPerCta; + static constexpr int kSplits = kGroupedVerifyWorkspaceRows / MAX_QUERY_TOKENS; + static_assert(MAX_QUERY_TOKENS * kHeadsPerCta == kGroupedVerifyRows, + "grouped verifier CTA must retain 48 rows"); + static_assert(kGroupedVerifyHeads % kHeadsPerCta == 0, + "query heads must divide evenly across grouped CTAs"); + static_assert(kSplits * MAX_QUERY_TOKENS == kGroupedVerifyWorkspaceRows, + "q8 and q16 workspaces must have equal element counts"); +}; + struct alignas(256) GroupedVerifySmem { union { struct { - alignas(16) __half q[kGroupedVerifyRows * kGroupedVerifyHeadDim]; - alignas(16) __half kv[kGroupedVerifyBlockN * kGroupedVerifyHeadDim]; - alignas(16) float scores[kGroupedVerifyRows * kGroupedVerifyBlockN]; - alignas(16) __half probs[kGroupedVerifyRows * kGroupedVerifyBlockN]; + alignas(16) __half q[kGroupedVerifyRows * kGroupedVerifyQStride]; + alignas(16) __half kv[kGroupedVerifyBlockN * kGroupedVerifyKVStride]; + alignas(16) float scores[kGroupedVerifyRows * kGroupedVerifyScoreStride]; + alignas(16) __half probs[kGroupedVerifyRows * kGroupedVerifyProbStride]; } compute; alignas(16) float output[kGroupedVerifyRows * kGroupedVerifyHeadDim]; } storage; alignas(16) float row_max[kGroupedVerifyRows]; alignas(16) float row_sum[kGroupedVerifyRows]; alignas(16) float row_scale[kGroupedVerifyRows]; + alignas(16) int page_ids[kGroupedVerifyPageIdsCapacity]; }; static_assert(sizeof(GroupedVerifySmem) <= 64 * 1024, @@ -1818,14 +1853,15 @@ static_assert(sizeof(GroupedVerifySmem) <= 64 * 1024, static_assert(kGroupedVerifyOutputTiles % kGroupedVerifyWarps == 0, "output tiles must divide evenly across warps"); -template +template __device__ __forceinline__ int grouped_verify_active_splits( const int total_kv) { + using Traits = GroupedVerifyTraits; constexpr int kMinTokensPerSplit = SINGLE_QUERY ? kGroupedVerifySingleQueryMinTokensPerSplit : kGroupedVerifyMinTokensPerSplit; const int active_splits = - min(kGroupedVerifySplits, + min(Traits::kSplits, max(1, (total_kv + kMinTokensPerSplit - 1) / kMinTokensPerSplit)); if constexpr (SINGLE_QUERY) { return active_splits; @@ -1853,13 +1889,13 @@ __device__ __forceinline__ void grouped_verify_qk( #pragma unroll for (int k_offset = 0; k_offset < kGroupedVerifyHeadDim; k_offset += 16) { volta::load_matrix_sync( - q_fragment, shared_q + m_tile * 16 * kGroupedVerifyHeadDim + k_offset, - kGroupedVerifyHeadDim); + q_fragment, shared_q + m_tile * 16 * kGroupedVerifyQStride + k_offset, + kGroupedVerifyQStride); // K is row-major [N, D]. The same bytes represent K^T as a col-major // [D, N] matrix, which is the B operand needed by Q @ K^T. volta::load_matrix_sync( - k_fragment, shared_k + n_tile * 16 * kGroupedVerifyHeadDim + k_offset, - kGroupedVerifyHeadDim); + k_fragment, shared_k + n_tile * 16 * kGroupedVerifyKVStride + k_offset, + kGroupedVerifyKVStride); volta::mma_sync(score_fragment, q_fragment, k_fragment, score_fragment); } #pragma unroll @@ -1867,8 +1903,8 @@ __device__ __forceinline__ void grouped_verify_qk( score_fragment.x[i] *= qk_scale; } volta::store_matrix_sync( - shared_scores + m_tile * 16 * kGroupedVerifyBlockN + n_tile * 16, - score_fragment, kGroupedVerifyBlockN, volta::mem_row_major); + shared_scores + m_tile * 16 * kGroupedVerifyScoreStride + n_tile * 16, + score_fragment, kGroupedVerifyScoreStride, volta::mem_row_major); } __device__ __forceinline__ void grouped_verify_scale_output_fragment( @@ -1888,7 +1924,9 @@ __device__ __forceinline__ void grouped_verify_scale_output_fragment( fragment.x[7] *= second_scale; } -template +template __global__ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_e5m2_partial_kernel( const __half* __restrict__ q, const void* __restrict__ k_cache, @@ -1900,11 +1938,12 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const int64_t k_head_stride, const int64_t v_block_stride, const int64_t v_token_stride, const int64_t v_head_stride, const float qk_scale, const float v_scale) { + using Traits = GroupedVerifyTraits; const int head_group = blockIdx.x; const int split_id = blockIdx.y; (void)max_num_blocks; - if (head_group >= 1 || split_id >= kGroupedVerifySplits || query_len <= 0 || - query_len > kGroupedVerifyMaxQ) { + if (head_group >= Traits::kHeadGroups || split_id >= Traits::kSplits || + query_len <= 0 || query_len > MAX_QUERY_TOKENS) { return; } @@ -1913,7 +1952,7 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ return; } const int active_splits = - grouped_verify_active_splits(total_kv); + grouped_verify_active_splits(total_kv); if (split_id >= active_splits) { return; } @@ -1928,7 +1967,7 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const int split_end = min(total_kv, split_start + split_tiles * kGroupedVerifyBlockN); const int prefix_kv_len = max(0, total_kv - query_len); - const int head_start = head_group * kGroupedVerifyHeadsPerCta; + const int head_start = head_group * Traits::kHeadsPerCta; const int tid = threadIdx.x; const int warp_id = tid / kWarpSize; const int lane_id = tid % kWarpSize; @@ -1940,23 +1979,46 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ __half* shared_kv = smem.storage.compute.kv; float* shared_scores = smem.storage.compute.scores; __half* shared_probs = smem.storage.compute.probs; + const int* page_ids = block_table; + int split_page_offset = 0; + bool use_staged_page_ids = false; + if constexpr (STAGE_PARTITION_PAGE_IDS) { + static_assert(PAGE_BLOCK_SIZE > 0, + "partition page staging requires a specialized page size"); + const int split_start_page = split_start / PAGE_BLOCK_SIZE; + split_page_offset = split_start - split_start_page * PAGE_BLOCK_SIZE; + const int split_page_count = + (split_page_offset + split_end - split_start + PAGE_BLOCK_SIZE - 1) / + PAGE_BLOCK_SIZE; + use_staged_page_ids = split_page_count <= kGroupedVerifyPageIdsCapacity; + if (use_staged_page_ids) { + for (int idx = tid; idx < split_page_count; + idx += kGroupedVerifyThreads) { + smem.page_ids[idx] = __ldg(&block_table[split_start_page + idx]); + } + } + if (use_staged_page_ids) { + page_ids = smem.page_ids; + } + } constexpr int kVecsPerRow = kGroupedVerifyHeadDim / 8; + constexpr int kSharedQVecsPerRow = kGroupedVerifyQStride / 8; const uint4* q_vec = reinterpret_cast(q); uint4* shared_q_vec = reinterpret_cast(shared_q); for (int idx = tid; idx < kGroupedVerifyRows * kVecsPerRow; idx += kGroupedVerifyThreads) { const int row = idx / kVecsPerRow; const int vec_col = idx % kVecsPerRow; - const int token_idx = row / kGroupedVerifyHeadsPerCta; - const int local_head = row % kGroupedVerifyHeadsPerCta; + const int token_idx = row / Traits::kHeadsPerCta; + const int local_head = row % Traits::kHeadsPerCta; const int head_idx = head_start + local_head; if (token_idx < query_len && head_idx < kGroupedVerifyHeads) { - shared_q_vec[idx] = __ldg( + shared_q_vec[row * kSharedQVecsPerRow + vec_col] = __ldg( q_vec + (token_idx * kGroupedVerifyHeads + head_idx) * kVecsPerRow + vec_col); } else { - shared_q_vec[idx] = make_uint4(0, 0, 0, 0); + shared_q_vec[row * kSharedQVecsPerRow + vec_col] = make_uint4(0, 0, 0, 0); } } if (tid < kGroupedVerifyRows) { @@ -1964,10 +2026,11 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ smem.row_sum[tid] = 0.0f; smem.row_scale[tid] = 1.0f; } + // This barrier publishes both Q and the optional split-local page IDs. __syncthreads(); constexpr int kPanelStrideVec = kGroupedVerifyHeadDim / 8; - constexpr int kSharedStrideVec = kGroupedVerifyHeadDim / 8; + constexpr int kSharedStrideVec = kGroupedVerifyKVStride / 8; volta::fragment output_fragments[kGroupedVerifyOutputTilesPerWarp]; @@ -1984,11 +2047,15 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ tile_start += kGroupedVerifyBlockN) { const int valid_k_rows = min(kGroupedVerifyBlockN, split_end - tile_start); - load_xqa_tc_kv_panel( - shared_kv, k_cache, block_table, valid_k_rows, kPanelStrideVec, - kSharedStrideVec, tile_start, 0, page_block_size, 0, k_block_stride, - k_token_stride, k_head_stride, 0); + shared_kv, k_cache, page_ids, valid_k_rows, kPanelStrideVec, + kSharedStrideVec, tile_page_offset, 0, page_block_size, 0, + k_block_stride, k_token_stride, k_head_stride, 0); for (int idx = tid + valid_k_rows * kSharedStrideVec; idx < kGroupedVerifyBlockN * kSharedStrideVec; idx += kGroupedVerifyThreads) { @@ -2002,15 +2069,15 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ #pragma unroll for (int row = warp_id; row < kGroupedVerifyRows; row += kGroupedVerifyWarps) { - const int token_idx = row / kGroupedVerifyHeadsPerCta; - const int local_head = row % kGroupedVerifyHeadsPerCta; + const int token_idx = row / Traits::kHeadsPerCta; + const int local_head = row % Traits::kHeadsPerCta; const int head_idx = head_start + local_head; const int kv_idx = tile_start + lane_id; const bool visible = token_idx < query_len && head_idx < kGroupedVerifyHeads && lane_id < valid_k_rows && kv_idx <= prefix_kv_len + token_idx; const float score = - visible ? shared_scores[row * kGroupedVerifyBlockN + lane_id] + visible ? shared_scores[row * kGroupedVerifyScoreStride + lane_id] : kXQANegInf; const float tile_max_lane = warp_reduce_max(score); const float tile_max = __shfl_sync(0xffffffffu, tile_max_lane, 0); @@ -2038,11 +2105,15 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ for (int tile_start = split_start; tile_start < split_end; tile_start += kGroupedVerifyBlockN) { const int valid_k_rows = min(kGroupedVerifyBlockN, split_end - tile_start); - load_xqa_tc_kv_panel( - shared_kv, k_cache, block_table, valid_k_rows, kPanelStrideVec, - kSharedStrideVec, tile_start, 0, page_block_size, 0, k_block_stride, - k_token_stride, k_head_stride, 0); + shared_kv, k_cache, page_ids, valid_k_rows, kPanelStrideVec, + kSharedStrideVec, tile_page_offset, 0, page_block_size, 0, + k_block_stride, k_token_stride, k_head_stride, 0); for (int idx = tid + valid_k_rows * kSharedStrideVec; idx < kGroupedVerifyBlockN * kSharedStrideVec; idx += kGroupedVerifyThreads) { @@ -2058,8 +2129,8 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ idx += kGroupedVerifyThreads) { const int row = idx / kGroupedVerifyBlockN; const int col = idx % kGroupedVerifyBlockN; - const int token_idx = row / kGroupedVerifyHeadsPerCta; - const int local_head = row % kGroupedVerifyHeadsPerCta; + const int token_idx = row / Traits::kHeadsPerCta; + const int local_head = row % Traits::kHeadsPerCta; const int head_idx = head_start + local_head; const int kv_idx = tile_start + col; const bool visible = @@ -2067,25 +2138,28 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ col < valid_k_rows && kv_idx <= prefix_kv_len + token_idx && smem.row_sum[row] > 0.0f; const float probability = - visible - ? __expf(fmaxf(shared_scores[idx] - smem.row_max[row], -80.0f)) - : 0.0f; - shared_probs[idx] = __float2half_rn(probability); + visible ? __expf(fmaxf( + shared_scores[row * kGroupedVerifyScoreStride + col] - + smem.row_max[row], + -80.0f)) + : 0.0f; + shared_probs[row * kGroupedVerifyProbStride + col] = + __float2half_rn(probability); } __syncthreads(); } else { #pragma unroll for (int row = warp_id; row < kGroupedVerifyRows; row += kGroupedVerifyWarps) { - const int token_idx = row / kGroupedVerifyHeadsPerCta; - const int local_head = row % kGroupedVerifyHeadsPerCta; + const int token_idx = row / Traits::kHeadsPerCta; + const int local_head = row % Traits::kHeadsPerCta; const int head_idx = head_start + local_head; const int kv_idx = tile_start + lane_id; const bool visible = token_idx < query_len && head_idx < kGroupedVerifyHeads && lane_id < valid_k_rows && kv_idx <= prefix_kv_len + token_idx; const float score = - visible ? shared_scores[row * kGroupedVerifyBlockN + lane_id] + visible ? shared_scores[row * kGroupedVerifyScoreStride + lane_id] : kXQANegInf; const float tile_max_lane = warp_reduce_max(score); const float tile_max = __shfl_sync(0xffffffffu, tile_max_lane, 0); @@ -2097,7 +2171,7 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ const float tile_sum = __shfl_sync(0xffffffffu, tile_sum_lane, 0); const float exp_diff = tile_sum > 0.0f ? __expf(fmaxf(old_max - new_max, -80.0f)) : 1.0f; - shared_probs[row * kGroupedVerifyBlockN + lane_id] = + shared_probs[row * kGroupedVerifyProbStride + lane_id] = __float2half_rn(probability); if (lane_id == 0) { if (tile_sum > 0.0f) { @@ -2118,11 +2192,12 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ } } - load_xqa_tc_kv_panel( - shared_kv, v_cache, block_table, valid_k_rows, kPanelStrideVec, - kSharedStrideVec, tile_start, 0, page_block_size, 0, v_block_stride, - v_token_stride, v_head_stride, 0); + shared_kv, v_cache, page_ids, valid_k_rows, kPanelStrideVec, + kSharedStrideVec, tile_page_offset, 0, page_block_size, 0, + v_block_stride, v_token_stride, v_head_stride, 0); for (int idx = tid + valid_k_rows * kSharedStrideVec; idx < kGroupedVerifyBlockN * kSharedStrideVec; idx += kGroupedVerifyThreads) { @@ -2144,12 +2219,12 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ for (int k_offset = 0; k_offset < kGroupedVerifyBlockN; k_offset += 16) { volta::load_matrix_sync( probability_fragment, - shared_probs + m_tile * 16 * kGroupedVerifyBlockN + k_offset, - kGroupedVerifyBlockN); + shared_probs + m_tile * 16 * kGroupedVerifyProbStride + k_offset, + kGroupedVerifyProbStride); volta::load_matrix_sync( value_fragment, - shared_kv + k_offset * kGroupedVerifyHeadDim + d_tile * 16, - kGroupedVerifyHeadDim); + shared_kv + k_offset * kGroupedVerifyKVStride + d_tile * 16, + kGroupedVerifyKVStride); volta::mma_sync(output_fragments[fragment_idx], probability_fragment, value_fragment, output_fragments[fragment_idx]); } @@ -2178,14 +2253,14 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ idx += kGroupedVerifyThreads) { const int row = idx / kGroupedVerifyHeadDim; const int d = idx % kGroupedVerifyHeadDim; - const int token_idx = row / kGroupedVerifyHeadsPerCta; - const int local_head = row % kGroupedVerifyHeadsPerCta; + const int token_idx = row / Traits::kHeadsPerCta; + const int local_head = row % Traits::kHeadsPerCta; const int head_idx = head_start + local_head; if (token_idx < query_len && head_idx < kGroupedVerifyHeads) { const float sum = smem.row_sum[row]; const float scale = sum > 0.0f ? v_scale / sum : 0.0f; const int64_t output_idx = - (((static_cast(split_id) * kGroupedVerifyMaxQ + token_idx) * + (((static_cast(split_id) * MAX_QUERY_TOKENS + token_idx) * kGroupedVerifyHeads + head_idx) * kGroupedVerifyHeadDim + @@ -2194,13 +2269,13 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ } } if (tid < kGroupedVerifyRows) { - const int token_idx = tid / kGroupedVerifyHeadsPerCta; - const int local_head = tid % kGroupedVerifyHeadsPerCta; + const int token_idx = tid / Traits::kHeadsPerCta; + const int local_head = tid % Traits::kHeadsPerCta; const int head_idx = head_start + local_head; if (token_idx < query_len && head_idx < kGroupedVerifyHeads) { const float sum = smem.row_sum[tid]; const int64_t lse_idx = - (static_cast(split_id) * kGroupedVerifyMaxQ + token_idx) * + (static_cast(split_id) * MAX_QUERY_TOKENS + token_idx) * kGroupedVerifyHeads + head_idx; partial_lse[lse_idx] = @@ -2209,26 +2284,27 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ } } -template +template __global__ __launch_bounds__(kGroupedVerifyThreads) void flash_attention_grouped_verify_e5m2_combine_kernel( const __half* __restrict__ partial_out, const float* __restrict__ partial_lse, const int* __restrict__ seq_lens, __half* __restrict__ out, const int query_len) { + using Traits = GroupedVerifyTraits; const int token_idx = blockIdx.x; const int head_idx = blockIdx.y; if (token_idx >= query_len || head_idx >= kGroupedVerifyHeads) { return; } const int active_splits = - grouped_verify_active_splits(seq_lens[0]); - __shared__ float split_lse[kGroupedVerifySplits]; + grouped_verify_active_splits(seq_lens[0]); + __shared__ float split_lse[Traits::kSplits]; __shared__ float final_max; __shared__ float final_inv_sum; - if (threadIdx.x < kGroupedVerifySplits) { + if (threadIdx.x < Traits::kSplits) { const int64_t lse_idx = - (static_cast(threadIdx.x) * kGroupedVerifyMaxQ + token_idx) * + (static_cast(threadIdx.x) * MAX_QUERY_TOKENS + token_idx) * kGroupedVerifyHeads + head_idx; split_lse[threadIdx.x] = @@ -2260,7 +2336,7 @@ __launch_bounds__(kGroupedVerifyThreads) void flash_attention_grouped_verify_e5m const float weight = __expf(fmaxf(split_lse[split] - final_max, -80.0f)) * final_inv_sum; const int64_t partial_idx = - (((static_cast(split) * kGroupedVerifyMaxQ + token_idx) * + (((static_cast(split) * MAX_QUERY_TOKENS + token_idx) * kGroupedVerifyHeads + head_idx) * kGroupedVerifyHeadDim + @@ -3463,10 +3539,16 @@ at::Tensor flash_attention_grouped_verify_paged( partial_lse.dtype() == torch::kFloat32, "grouped verify workspaces must be fp16/fp32"); TORCH_CHECK(q.dim() == 3 && q.size(0) > 0 && - q.size(0) <= kGroupedVerifyMaxQ && + q.size(0) <= kGroupedVerifyMaxSupportedQ && q.size(1) == kGroupedVerifyHeads && q.size(2) == kGroupedVerifyHeadDim, - "grouped verify q must have shape [1..8, 6, 256]"); + "grouped verify q must have shape [1..16, 6, 256]"); + const bool wide_query = q.size(0) > kGroupedVerifyQ8MaxQ; + const int max_query_tokens = + wide_query ? kGroupedVerifyQ16MaxQ : kGroupedVerifyQ8MaxQ; + const int grouped_splits = kGroupedVerifyWorkspaceRows / max_query_tokens; + const int heads_per_cta = kGroupedVerifyRows / max_query_tokens; + const int head_groups = kGroupedVerifyHeads / heads_per_cta; TORCH_CHECK(k_cache.dim() == 4 && v_cache.dim() == 4 && k_cache.sizes() == v_cache.sizes() && k_cache.size(1) > 0 && k_cache.size(2) == 1 && @@ -3485,13 +3567,14 @@ at::Tensor flash_attention_grouped_verify_paged( TORCH_CHECK(partial_out.is_contiguous() && partial_lse.is_contiguous(), "grouped verify workspaces must be contiguous"); TORCH_CHECK(partial_out.sizes() == - at::IntArrayRef({kGroupedVerifySplits, kGroupedVerifyMaxQ, + at::IntArrayRef({grouped_splits, max_query_tokens, kGroupedVerifyHeads, kGroupedVerifyHeadDim}), - "partial_out must have shape [80, 8, 6, 256]"); - TORCH_CHECK(partial_lse.sizes() == - at::IntArrayRef({kGroupedVerifySplits, kGroupedVerifyMaxQ, - kGroupedVerifyHeads}), - "partial_lse must have shape [80, 8, 6]"); + "partial_out must have shape [80, 8, 6, 256] or " + "[40, 16, 6, 256]"); + TORCH_CHECK( + partial_lse.sizes() == at::IntArrayRef({grouped_splits, max_query_tokens, + kGroupedVerifyHeads}), + "partial_lse must have shape [80, 8, 6] or [40, 16, 6]"); TORCH_CHECK(k_scale > 0.0f && v_scale > 0.0f, "grouped verify E5M2 K/V scales must be positive"); @@ -3514,13 +3597,16 @@ at::Tensor flash_attention_grouped_verify_paged( "grouped verify prototype supports SM70 only"); cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - const dim3 partial_grid(1, kGroupedVerifySplits, 1); + const dim3 partial_grid(head_groups, grouped_splits, 1); const size_t partial_shared_mem = sizeof(GroupedVerifySmem); -#define LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, PAGE_SIZE, SINGLE_QUERY) \ +#define LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, PAGE_SIZE, \ + SINGLE_QUERY, CONTIGUOUS_LAYOUT, \ + STAGE_PAGE_IDS) \ do { \ auto partial_kernel = \ (void*)flash_attention_grouped_verify_e5m2_partial_kernel< \ - TWO_PASS, PAGE_SIZE, SINGLE_QUERY>; \ + MAX_QUERY_TOKENS, TWO_PASS, PAGE_SIZE, SINGLE_QUERY, \ + CONTIGUOUS_LAYOUT, STAGE_PAGE_IDS>; \ const cudaError_t smem_status = cudaFuncSetAttribute( \ partial_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, \ partial_shared_mem); \ @@ -3532,8 +3618,9 @@ at::Tensor flash_attention_grouped_verify_paged( TORCH_CHECK(carveout_status == cudaSuccess, \ "Failed to set packed grouped-verifier shared carveout: ", \ cudaGetErrorString(carveout_status)); \ - flash_attention_grouped_verify_e5m2_partial_kernel \ + flash_attention_grouped_verify_e5m2_partial_kernel< \ + MAX_QUERY_TOKENS, TWO_PASS, PAGE_SIZE, SINGLE_QUERY, \ + CONTIGUOUS_LAYOUT, STAGE_PAGE_IDS> \ <<>>( \ reinterpret_cast(q.data_ptr()), k_cache.data_ptr(), \ v_cache.data_ptr(), block_table.data_ptr(), \ @@ -3547,42 +3634,78 @@ at::Tensor flash_attention_grouped_verify_paged( v_scale); \ } while (0) -#define DISPATCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, SINGLE_QUERY) \ - do { \ - if (k_cache.size(1) == 1648) { \ - LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 1648, SINGLE_QUERY); \ - } else if (k_cache.size(1) == 3296) { \ - LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 3296, SINGLE_QUERY); \ - } else { \ - LAUNCH_GROUPED_VERIFY_PARTIAL(TWO_PASS, 0, SINGLE_QUERY); \ - } \ +#define DISPATCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, \ + SINGLE_QUERY) \ + do { \ + const int64_t page_size = k_cache.size(1); \ + const int64_t fixed_block_stride = 2 * page_size * kGroupedVerifyHeadDim; \ + const bool fixed_interleaved_layout = \ + dflash2_grouped_fixed_interleaved_enabled() && \ + (page_size == 1648 || page_size == 3296) && \ + k_cache.stride(0) == fixed_block_stride && \ + v_cache.stride(0) == fixed_block_stride && \ + k_cache.stride(1) == kGroupedVerifyHeadDim && \ + v_cache.stride(1) == kGroupedVerifyHeadDim && \ + k_cache.stride(2) == kGroupedVerifyHeadDim && \ + v_cache.stride(2) == kGroupedVerifyHeadDim; \ + const bool stage_page_ids = \ + fixed_interleaved_layout && dflash2_grouped_stage_page_ids_enabled(); \ + if (stage_page_ids && page_size == 1648) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 1648, \ + SINGLE_QUERY, true, true); \ + } else if (stage_page_ids) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 3296, \ + SINGLE_QUERY, true, true); \ + } else if (fixed_interleaved_layout && page_size == 1648) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 1648, \ + SINGLE_QUERY, true, false); \ + } else if (fixed_interleaved_layout) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 3296, \ + SINGLE_QUERY, true, false); \ + } else if (page_size == 1648) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 1648, \ + SINGLE_QUERY, false, false); \ + } else if (page_size == 3296) { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 3296, \ + SINGLE_QUERY, false, false); \ + } else { \ + LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 0, \ + SINGLE_QUERY, false, false); \ + } \ } while (0) const bool single_query = q.size(0) == 1; - if (one_pass && single_query) { - DISPATCH_GROUPED_VERIFY_PARTIAL(false, true); + if (wide_query && one_pass) { + DISPATCH_GROUPED_VERIFY_PARTIAL(kGroupedVerifyQ16MaxQ, false, false); + } else if (wide_query) { + DISPATCH_GROUPED_VERIFY_PARTIAL(kGroupedVerifyQ16MaxQ, true, false); + } else if (one_pass && single_query) { + DISPATCH_GROUPED_VERIFY_PARTIAL(kGroupedVerifyQ8MaxQ, false, true); } else if (one_pass) { - DISPATCH_GROUPED_VERIFY_PARTIAL(false, false); + DISPATCH_GROUPED_VERIFY_PARTIAL(kGroupedVerifyQ8MaxQ, false, false); } else if (single_query) { - DISPATCH_GROUPED_VERIFY_PARTIAL(true, true); + DISPATCH_GROUPED_VERIFY_PARTIAL(kGroupedVerifyQ8MaxQ, true, true); } else { - DISPATCH_GROUPED_VERIFY_PARTIAL(true, false); + DISPATCH_GROUPED_VERIFY_PARTIAL(kGroupedVerifyQ8MaxQ, true, false); } #undef DISPATCH_GROUPED_VERIFY_PARTIAL #undef LAUNCH_GROUPED_VERIFY_PARTIAL const dim3 combine_grid(static_cast(q.size(0)), kGroupedVerifyHeads, 1); -#define LAUNCH_GROUPED_VERIFY_COMBINE(SINGLE_QUERY) \ - flash_attention_grouped_verify_e5m2_combine_kernel \ - <<>>( \ - reinterpret_cast(partial_out.data_ptr()), \ - partial_lse.data_ptr(), seq_lens.data_ptr(), \ - reinterpret_cast<__half*>(out.data_ptr()), \ +#define LAUNCH_GROUPED_VERIFY_COMBINE(MAX_QUERY_TOKENS, SINGLE_QUERY) \ + flash_attention_grouped_verify_e5m2_combine_kernel \ + <<>>( \ + reinterpret_cast(partial_out.data_ptr()), \ + partial_lse.data_ptr(), seq_lens.data_ptr(), \ + reinterpret_cast<__half*>(out.data_ptr()), \ static_cast(q.size(0))) - if (single_query) { - LAUNCH_GROUPED_VERIFY_COMBINE(true); + if (wide_query) { + LAUNCH_GROUPED_VERIFY_COMBINE(kGroupedVerifyQ16MaxQ, false); + } else if (single_query) { + LAUNCH_GROUPED_VERIFY_COMBINE(kGroupedVerifyQ8MaxQ, true); } else { - LAUNCH_GROUPED_VERIFY_COMBINE(false); + LAUNCH_GROUPED_VERIFY_COMBINE(kGroupedVerifyQ8MaxQ, false); } #undef LAUNCH_GROUPED_VERIFY_COMBINE C10_CUDA_KERNEL_LAUNCH_CHECK(); diff --git a/flash-attention-v100/kernel/flash_decode_turboquant.cu b/flash-attention-v100/kernel/flash_decode_turboquant.cu index 79c7630612..41627bd78d 100644 --- a/flash-attention-v100/kernel/flash_decode_turboquant.cu +++ b/flash-attention-v100/kernel/flash_decode_turboquant.cu @@ -14,7 +14,7 @@ constexpr int kThreadsPerBlock = 256; constexpr int kWarpsPerBlock = kThreadsPerBlock / kWarpSize; __device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll +#pragma unroll for (int offset = kWarpSize / 2; offset > 0; offset /= 2) { val += __shfl_down_sync(0xffffffff, val, offset); } @@ -22,14 +22,14 @@ __device__ __forceinline__ float warp_reduce_sum(float val) { } __device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll +#pragma unroll for (int offset = kWarpSize / 2; offset > 0; offset /= 2) { val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset)); } return val; } -template +template __device__ __forceinline__ float block_reduce_sum(float val) { __shared__ float shared[NUM_WARPS]; __shared__ float result; @@ -53,7 +53,7 @@ __device__ __forceinline__ float block_reduce_sum(float val) { return result; } -template +template __device__ __forceinline__ float block_reduce_max(float val) { __shared__ float shared[NUM_WARPS]; __shared__ float result; @@ -77,32 +77,28 @@ __device__ __forceinline__ float block_reduce_max(float val) { return result; } -__device__ __forceinline__ uint16_t load_u16_le( - const uint8_t* __restrict__ data, - const int64_t byte_index) { +__device__ __forceinline__ uint16_t +load_u16_le(const uint8_t* __restrict__ data, const int64_t byte_index) { return static_cast(data[byte_index]) | (static_cast(data[byte_index + 1]) << 8); } __device__ __forceinline__ float load_f16_le_float( - const uint8_t* __restrict__ data, - const int64_t byte_index) { + const uint8_t* __restrict__ data, const int64_t byte_index) { return __half2float(__ushort_as_half(load_u16_le(data, byte_index))); } -template +template __device__ __forceinline__ float dot_qk_turboquant_mse( - const float* __restrict__ q_rot, - const uint8_t* __restrict__ kv_cache, - const int64_t slot_base, - const float* __restrict__ centroids, + const float* __restrict__ q_rot, const uint8_t* __restrict__ kv_cache, + const int64_t slot_base, const float* __restrict__ centroids, const int lane) { constexpr int kMseBytes = (D * MSE_BITS + 7) / 8; constexpr int kMseMask = (1 << MSE_BITS) - 1; float acc = 0.f; float centroid_norm_sq = 0.f; - #pragma unroll +#pragma unroll for (int d = lane; d < D; d += kWarpSize) { const int bit_offset = d * MSE_BITS; const int byte_offset = bit_offset >> 3; @@ -126,10 +122,9 @@ __device__ __forceinline__ float dot_qk_turboquant_mse( return acc * vec_norm; } -template +template __device__ __forceinline__ float load_turboquant_value_float( - const uint8_t* __restrict__ kv_cache, - const int64_t slot_base, + const uint8_t* __restrict__ kv_cache, const int64_t slot_base, const int d) { constexpr int kMseBytes = (D * MSE_BITS + 7) / 8; constexpr int kKeyPackedSize = kMseBytes + 2; @@ -154,32 +149,20 @@ __device__ __forceinline__ float load_turboquant_value_float( return static_cast(q_value) * scale + zero; } -template +template __global__ void flash_attention_turboquant_decode_partition_kernel( - const float* __restrict__ q_rot, - const uint8_t* __restrict__ kv_cache, - float* __restrict__ tmp_out, - float* __restrict__ max_logits, - float* __restrict__ exp_sums, - const int* __restrict__ block_table, - const int* __restrict__ seq_lens, - const float* __restrict__ centroids, - const int batch_size, - const int max_num_blocks, - const int max_num_partitions, - const int num_heads_q, - const int num_heads_kv, - const int block_size, - const int64_t q_stride0, - const int64_t q_stride1, - const int64_t tmp_out_stride0, - const int64_t tmp_out_stride1, - const int64_t tmp_out_stride2, - const int64_t stats_stride0, - const int64_t stats_stride1, - const int64_t cache_block_stride, - const int64_t cache_token_stride, - const int64_t cache_head_stride, + const float* __restrict__ q_rot, const uint8_t* __restrict__ kv_cache, + float* __restrict__ tmp_out, float* __restrict__ max_logits, + float* __restrict__ exp_sums, const int* __restrict__ block_table, + const int* __restrict__ seq_lens, const float* __restrict__ centroids, + const int batch_size, const int max_num_blocks, + const int max_num_partitions, const int num_heads_q, const int num_heads_kv, + const int block_size, const int64_t q_stride0, const int64_t q_stride1, + const int64_t tmp_out_stride0, const int64_t tmp_out_stride1, + const int64_t tmp_out_stride2, const int64_t stats_stride0, + const int64_t stats_stride1, const int64_t cache_block_stride, + const int64_t cache_token_stride, const int64_t cache_head_stride, const float softmax_scale) { const int batch_idx = blockIdx.x; const int head_idx = blockIdx.y; @@ -191,8 +174,7 @@ __global__ void flash_attention_turboquant_decode_partition_kernel( } const int seq_len = seq_lens[batch_idx]; - const int split_len = - (seq_len + max_num_partitions - 1) / max_num_partitions; + const int split_len = (seq_len + max_num_partitions - 1) / max_num_partitions; const int start_token_idx = partition_idx * split_len; if (seq_len <= 0 || start_token_idx >= seq_len) { return; @@ -268,8 +250,8 @@ __global__ void flash_attention_turboquant_decode_partition_kernel( static_cast(physical_block) * cache_block_stride + static_cast(block_offset) * cache_token_stride + static_cast(kv_head_idx) * cache_head_stride; - const float value = load_turboquant_value_float( - kv_cache, slot_base, d); + const float value = + load_turboquant_value_float(kv_cache, slot_base, d); acc = fmaf(scores_shared[i], value, acc); } tmp_out[tmp_out_base + d] = acc * inv_part_sum; @@ -284,22 +266,15 @@ __global__ void flash_attention_turboquant_decode_partition_kernel( } } -template +template __global__ void flash_attention_turboquant_decode_reduce_kernel( - const float* __restrict__ tmp_out, - const float* __restrict__ max_logits, - const float* __restrict__ exp_sums, - const int* __restrict__ seq_lens, - __half* __restrict__ out, - const int batch_size, - const int max_num_partitions, - const int num_heads_q, - const int64_t tmp_out_stride0, - const int64_t tmp_out_stride1, - const int64_t tmp_out_stride2, - const int64_t stats_stride0, - const int64_t stats_stride1, - const int64_t out_stride0, + const float* __restrict__ tmp_out, const float* __restrict__ max_logits, + const float* __restrict__ exp_sums, const int* __restrict__ seq_lens, + __half* __restrict__ out, const int batch_size, + const int max_num_partitions, const int num_heads_q, + const int64_t tmp_out_stride0, const int64_t tmp_out_stride1, + const int64_t tmp_out_stride2, const int64_t stats_stride0, + const int64_t stats_stride1, const int64_t out_stride0, const int64_t out_stride1) { const int batch_idx = blockIdx.x; const int head_idx = blockIdx.y; @@ -313,13 +288,11 @@ __global__ void flash_attention_turboquant_decode_reduce_kernel( if (seq_len <= 0) { for (int d = threadIdx.x; d < D; d += blockDim.x) { out[static_cast(batch_idx) * out_stride0 + - static_cast(head_idx) * out_stride1 + d] = - __float2half(0.f); + static_cast(head_idx) * out_stride1 + d] = __float2half(0.f); } return; } - const int split_len = - (seq_len + max_num_partitions - 1) / max_num_partitions; + const int split_len = (seq_len + max_num_partitions - 1) / max_num_partitions; const int num_partitions = (seq_len + split_len - 1) / split_len; extern __shared__ float shared_mem[]; @@ -342,7 +315,8 @@ __global__ void flash_attention_turboquant_decode_reduce_kernel( const int64_t stats_index = static_cast(batch_idx) * stats_stride0 + static_cast(head_idx) * stats_stride1 + i; - const float weight = exp_sums[stats_index] * __expf(max_shared[i] - global_max); + const float weight = + exp_sums[stats_index] * __expf(max_shared[i] - global_max); weight_shared[i] = weight; local_sum += weight; } @@ -350,9 +324,8 @@ __global__ void flash_attention_turboquant_decode_reduce_kernel( const float inv_global_sum = global_sum > 0.f ? 1.f / global_sum : 0.f; __syncthreads(); - const int64_t out_base = - static_cast(batch_idx) * out_stride0 + - static_cast(head_idx) * out_stride1; + const int64_t out_base = static_cast(batch_idx) * out_stride0 + + static_cast(head_idx) * out_stride1; const int64_t tmp_out_base = static_cast(batch_idx) * tmp_out_stride0 + static_cast(head_idx) * tmp_out_stride1; @@ -369,18 +342,13 @@ __global__ void flash_attention_turboquant_decode_reduce_kernel( } } -template +template void launch_flash_attention_turboquant_decode_paged( - const at::Tensor& q_rot, - const at::Tensor& kv_cache, - at::Tensor& out, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - at::Tensor& tmp_out, - at::Tensor& max_logits, - at::Tensor& exp_sums, - const at::Tensor& centroids, - const float softmax_scale, + const at::Tensor& q_rot, const at::Tensor& kv_cache, at::Tensor& out, + const at::Tensor& block_table, const at::Tensor& seq_lens, + at::Tensor& tmp_out, at::Tensor& max_logits, at::Tensor& exp_sums, + const at::Tensor& centroids, const float softmax_scale, cudaStream_t stream) { const int batch_size = q_rot.size(0); const int num_heads_q = q_rot.size(1); @@ -398,68 +366,35 @@ void launch_flash_attention_turboquant_decode_paged( flash_attention_turboquant_decode_partition_kernel< D, PARTITION_SIZE, MSE_BITS, VQB, NORM_CORRECTION> <<>>( - q_rot.data_ptr(), - kv_cache.data_ptr(), - tmp_out.data_ptr(), - max_logits.data_ptr(), - exp_sums.data_ptr(), - block_table.data_ptr(), - seq_lens.data_ptr(), - centroids.data_ptr(), - batch_size, - max_num_blocks, - max_num_partitions, - num_heads_q, - num_heads_kv, - block_size, - q_rot.stride(0), - q_rot.stride(1), - tmp_out.stride(0), - tmp_out.stride(1), - tmp_out.stride(2), - max_logits.stride(0), - max_logits.stride(1), - kv_cache.stride(0), - kv_cache.stride(1), - kv_cache.stride(2), - softmax_scale); + q_rot.data_ptr(), kv_cache.data_ptr(), + tmp_out.data_ptr(), max_logits.data_ptr(), + exp_sums.data_ptr(), block_table.data_ptr(), + seq_lens.data_ptr(), centroids.data_ptr(), batch_size, + max_num_blocks, max_num_partitions, num_heads_q, num_heads_kv, + block_size, q_rot.stride(0), q_rot.stride(1), tmp_out.stride(0), + tmp_out.stride(1), tmp_out.stride(2), max_logits.stride(0), + max_logits.stride(1), kv_cache.stride(0), kv_cache.stride(1), + kv_cache.stride(2), softmax_scale); flash_attention_turboquant_decode_reduce_kernel <<>>( - tmp_out.data_ptr(), - max_logits.data_ptr(), - exp_sums.data_ptr(), - seq_lens.data_ptr(), - reinterpret_cast<__half*>(out.data_ptr()), - batch_size, - max_num_partitions, - num_heads_q, - tmp_out.stride(0), - tmp_out.stride(1), - tmp_out.stride(2), - max_logits.stride(0), - max_logits.stride(1), - out.stride(0), - out.stride(1)); + tmp_out.data_ptr(), max_logits.data_ptr(), + exp_sums.data_ptr(), seq_lens.data_ptr(), + reinterpret_cast<__half*>(out.data_ptr()), batch_size, + max_num_partitions, num_heads_q, tmp_out.stride(0), tmp_out.stride(1), + tmp_out.stride(2), max_logits.stride(0), max_logits.stride(1), + out.stride(0), out.stride(1)); } } // namespace at::Tensor flash_attention_turboquant_decode_paged( - const at::Tensor& q_rot, - const at::Tensor& kv_cache, - std::optional& out_, - const at::Tensor& block_table, - const at::Tensor& seq_lens, - at::Tensor& tmp_out, - at::Tensor& max_logits, - at::Tensor& exp_sums, - const at::Tensor& centroids, - const float softmax_scale, - const int partition_size, - const int mse_bits, - const int value_quant_bits, - const bool norm_correction) { + const at::Tensor& q_rot, const at::Tensor& kv_cache, + std::optional& out_, const at::Tensor& block_table, + const at::Tensor& seq_lens, at::Tensor& tmp_out, at::Tensor& max_logits, + at::Tensor& exp_sums, const at::Tensor& centroids, + const float softmax_scale, const int partition_size, const int mse_bits, + const int value_quant_bits, const bool norm_correction) { TORCH_CHECK(q_rot.is_cuda(), "q_rot must be on CUDA"); TORCH_CHECK(kv_cache.is_cuda(), "kv_cache must be on CUDA"); TORCH_CHECK(block_table.is_cuda() && seq_lens.is_cuda(), @@ -474,7 +409,8 @@ at::Tensor flash_attention_turboquant_decode_paged( TORCH_CHECK(tmp_out.dtype() == torch::kFloat32, "tmp_out must be fp32"); TORCH_CHECK(max_logits.dtype() == torch::kFloat32, "max_logits must be fp32"); TORCH_CHECK(exp_sums.dtype() == torch::kFloat32, "exp_sums must be fp32"); - TORCH_CHECK(block_table.dtype() == torch::kInt32, "block_table must be int32"); + TORCH_CHECK(block_table.dtype() == torch::kInt32, + "block_table must be int32"); TORCH_CHECK(seq_lens.dtype() == torch::kInt32, "seq_lens must be int32"); TORCH_CHECK(q_rot.dim() == 3, "q_rot must have shape [B, H, D]"); TORCH_CHECK(kv_cache.dim() == 4, @@ -483,7 +419,8 @@ at::Tensor flash_attention_turboquant_decode_paged( "block_table must have shape [B, max_num_blocks]"); TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must have shape [B]"); TORCH_CHECK(tmp_out.dim() == 4, "tmp_out must have shape [B_cap, H, P, D]"); - TORCH_CHECK(max_logits.dim() == 3, "max_logits must have shape [B_cap, H, P]"); + TORCH_CHECK(max_logits.dim() == 3, + "max_logits must have shape [B_cap, H, P]"); TORCH_CHECK(exp_sums.dim() == 3, "exp_sums must have shape [B_cap, H, P]"); TORCH_CHECK(q_rot.stride(-1) == 1, "q_rot last dim must be contiguous"); TORCH_CHECK(kv_cache.stride(-1) == 1, "kv_cache last dim must be contiguous"); @@ -500,23 +437,26 @@ at::Tensor flash_attention_turboquant_decode_paged( "seq_lens batch size must cover q_rot batch size"); TORCH_CHECK(q_rot.size(0) <= tmp_out.size(0), "tmp_out batch size must cover q_rot batch size"); - TORCH_CHECK(num_heads_q == tmp_out.size(1), "tmp_out head dimension mismatch"); + TORCH_CHECK(num_heads_q == tmp_out.size(1), + "tmp_out head dimension mismatch"); TORCH_CHECK(head_dim == tmp_out.size(3), "tmp_out head_dim mismatch"); TORCH_CHECK(max_logits.size(0) == tmp_out.size(0) && max_logits.size(1) == tmp_out.size(1) && max_logits.size(2) == tmp_out.size(2), "max_logits shape mismatch"); - TORCH_CHECK(exp_sums.sizes() == max_logits.sizes(), "exp_sums shape mismatch"); + TORCH_CHECK(exp_sums.sizes() == max_logits.sizes(), + "exp_sums shape mismatch"); TORCH_CHECK(num_heads_q % num_heads_kv == 0, "num_heads_q must be divisible by num_heads_kv"); - TORCH_CHECK(partition_size == 256 || partition_size == 512 || - partition_size == 1024, - "Unsupported decode partition_size: ", partition_size); TORCH_CHECK( - block_table.size(1) * kv_cache.size(1) <= partition_size * tmp_out.size(2), - "TurboQuant decode workspace cannot cover block_table capacity"); - TORCH_CHECK(mse_bits == 3 || mse_bits == 4, - "TurboQuant Flash-V100 decode currently supports 3/4-bit MSE keys"); + partition_size == 256 || partition_size == 512 || partition_size == 1024, + "Unsupported decode partition_size: ", partition_size); + TORCH_CHECK(block_table.size(1) * kv_cache.size(1) <= + partition_size * tmp_out.size(2), + "TurboQuant decode workspace cannot cover block_table capacity"); + TORCH_CHECK( + mse_bits == 3 || mse_bits == 4, + "TurboQuant Flash-V100 decode currently supports 3/4-bit MSE keys"); TORCH_CHECK(value_quant_bits == 3 || value_quant_bits == 4, "TurboQuant Flash-V100 decode currently supports 3/4-bit values"); TORCH_CHECK(centroids.numel() >= (1 << mse_bits), @@ -537,65 +477,67 @@ at::Tensor flash_attention_turboquant_decode_paged( auto stream = at::cuda::getCurrentCUDAStream().stream(); c10::cuda::CUDAGuard device_guard(q_rot.device()); - #define LAUNCH_TYPED(HDIM, PARTITION, MSE, VBITS, NC) \ - launch_flash_attention_turboquant_decode_paged( \ - q_rot, kv_cache, out, block_table, seq_lens, tmp_out, max_logits, \ - exp_sums, centroids, softmax_scale, stream) - - #define LAUNCH_BY_NORM(HDIM, PARTITION, MSE, VBITS) \ - do { \ - if (norm_correction) { \ - LAUNCH_TYPED(HDIM, PARTITION, MSE, VBITS, true); \ - } else { \ - LAUNCH_TYPED(HDIM, PARTITION, MSE, VBITS, false); \ - } \ - } while (0) - - #define LAUNCH_BY_VALUE(HDIM, PARTITION, MSE) \ - do { \ - switch (value_quant_bits) { \ - case 3: \ - LAUNCH_BY_NORM(HDIM, PARTITION, MSE, 3); \ - break; \ - case 4: \ - LAUNCH_BY_NORM(HDIM, PARTITION, MSE, 4); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported value_quant_bits: ", value_quant_bits); \ - } \ - } while (0) - - #define LAUNCH_BY_MSE(HDIM, PARTITION) \ - do { \ - switch (mse_bits) { \ - case 3: \ - LAUNCH_BY_VALUE(HDIM, PARTITION, 3); \ - break; \ - case 4: \ - LAUNCH_BY_VALUE(HDIM, PARTITION, 4); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported mse_bits: ", mse_bits); \ - } \ - } while (0) - - #define LAUNCH_BY_PARTITION(HDIM) \ - do { \ - switch (partition_size) { \ - case 256: \ - LAUNCH_BY_MSE(HDIM, 256); \ - break; \ - case 512: \ - LAUNCH_BY_MSE(HDIM, 512); \ - break; \ - case 1024: \ - LAUNCH_BY_MSE(HDIM, 1024); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported decode partition_size: ", partition_size); \ - } \ - } while (0) +#define LAUNCH_TYPED(HDIM, PARTITION, MSE, VBITS, NC) \ + launch_flash_attention_turboquant_decode_paged( \ + q_rot, kv_cache, out, block_table, seq_lens, tmp_out, max_logits, \ + exp_sums, centroids, softmax_scale, stream) + +#define LAUNCH_BY_NORM(HDIM, PARTITION, MSE, VBITS) \ + do { \ + if (norm_correction) { \ + LAUNCH_TYPED(HDIM, PARTITION, MSE, VBITS, true); \ + } else { \ + LAUNCH_TYPED(HDIM, PARTITION, MSE, VBITS, false); \ + } \ + } while (0) + +#define LAUNCH_BY_VALUE(HDIM, PARTITION, MSE) \ + do { \ + switch (value_quant_bits) { \ + case 3: \ + LAUNCH_BY_NORM(HDIM, PARTITION, MSE, 3); \ + break; \ + case 4: \ + LAUNCH_BY_NORM(HDIM, PARTITION, MSE, 4); \ + break; \ + default: \ + TORCH_CHECK(false, \ + "Unsupported value_quant_bits: ", value_quant_bits); \ + } \ + } while (0) + +#define LAUNCH_BY_MSE(HDIM, PARTITION) \ + do { \ + switch (mse_bits) { \ + case 3: \ + LAUNCH_BY_VALUE(HDIM, PARTITION, 3); \ + break; \ + case 4: \ + LAUNCH_BY_VALUE(HDIM, PARTITION, 4); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported mse_bits: ", mse_bits); \ + } \ + } while (0) + +#define LAUNCH_BY_PARTITION(HDIM) \ + do { \ + switch (partition_size) { \ + case 256: \ + LAUNCH_BY_MSE(HDIM, 256); \ + break; \ + case 512: \ + LAUNCH_BY_MSE(HDIM, 512); \ + break; \ + case 1024: \ + LAUNCH_BY_MSE(HDIM, 1024); \ + break; \ + default: \ + TORCH_CHECK(false, \ + "Unsupported decode partition_size: ", partition_size); \ + } \ + } while (0) switch (head_dim) { case 64: @@ -617,14 +559,15 @@ at::Tensor flash_attention_turboquant_decode_paged( LAUNCH_BY_PARTITION(256); break; default: - TORCH_CHECK(false, "Unsupported head_dim for TurboQuant paged decode: ", head_dim); + TORCH_CHECK(false, "Unsupported head_dim for TurboQuant paged decode: ", + head_dim); } - #undef LAUNCH_BY_PARTITION - #undef LAUNCH_BY_MSE - #undef LAUNCH_BY_VALUE - #undef LAUNCH_BY_NORM - #undef LAUNCH_TYPED +#undef LAUNCH_BY_PARTITION +#undef LAUNCH_BY_MSE +#undef LAUNCH_BY_VALUE +#undef LAUNCH_BY_NORM +#undef LAUNCH_TYPED C10_CUDA_KERNEL_LAUNCH_CHECK(); return out; diff --git a/flash-attention-v100/kernel/flash_v100_traits.cuh b/flash-attention-v100/kernel/flash_v100_traits.cuh index 909ea97db4..eaa191e2c3 100644 --- a/flash-attention-v100/kernel/flash_v100_traits.cuh +++ b/flash-attention-v100/kernel/flash_v100_traits.cuh @@ -3,137 +3,116 @@ #include #include -template +template struct FlashV100Traits { - - static constexpr int BLOCK_M_16 = 16; - static constexpr int BLOCK_N_16 = 512; - static constexpr int BLOCK_M_32 = 32; - static constexpr int BLOCK_N_32 = 256; - static constexpr int BLOCK_M_64 = 64; - static constexpr int BLOCK_N_64 = 128; - static constexpr int BLOCK_M_128 = 32; - static constexpr int BLOCK_N_128 = 176; - static constexpr int BLOCK_M_256 = 32; - static constexpr int BLOCK_N_256 = 64; - static constexpr int WARPS_PER_BLOCK = 16; - static constexpr int THREADS_PER_WARP = 32; - - static constexpr int BLOCK_M = (D == 16) ? BLOCK_M_16 : - (D == 32) ? BLOCK_M_32 : - (D == 64) ? BLOCK_M_64 : - (D == 128) ? BLOCK_M_128 : BLOCK_M_256; - - static constexpr int BLOCK_N = (D == 16) ? BLOCK_N_16 : - (D == 32) ? BLOCK_N_32 : - (D == 64) ? BLOCK_N_64 : - (D == 128) ? BLOCK_N_128 : BLOCK_N_256; - - static constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * THREADS_PER_WARP; - static constexpr int THREADS_PER_ROW = THREADS_PER_BLOCK / BLOCK_M; - - static constexpr int kBlockM = BLOCK_M; - static constexpr int kBlockN = BLOCK_N; - static constexpr int kHeadDim = D; - - static constexpr int kGmemThreadsPerRow = THREADS_PER_ROW; - static constexpr int kGmemRowsPerThread = 1; - static constexpr int kGmemElemsPerLoad = 8; - - static constexpr int PER_UINT4 = 8; - static constexpr int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; + static constexpr int BLOCK_M_16 = 16; + static constexpr int BLOCK_N_16 = 512; + static constexpr int BLOCK_M_32 = 32; + static constexpr int BLOCK_N_32 = 256; + static constexpr int BLOCK_M_64 = 64; + static constexpr int BLOCK_N_64 = 128; + static constexpr int BLOCK_M_128 = 32; + static constexpr int BLOCK_N_128 = 176; + static constexpr int BLOCK_M_256 = 32; + static constexpr int BLOCK_N_256 = 64; + static constexpr int WARPS_PER_BLOCK = 16; + static constexpr int THREADS_PER_WARP = 32; + + static constexpr int BLOCK_M = (D == 16) ? BLOCK_M_16 + : (D == 32) ? BLOCK_M_32 + : (D == 64) ? BLOCK_M_64 + : (D == 128) ? BLOCK_M_128 + : BLOCK_M_256; + + static constexpr int BLOCK_N = (D == 16) ? BLOCK_N_16 + : (D == 32) ? BLOCK_N_32 + : (D == 64) ? BLOCK_N_64 + : (D == 128) ? BLOCK_N_128 + : BLOCK_N_256; + + static constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * THREADS_PER_WARP; + static constexpr int THREADS_PER_ROW = THREADS_PER_BLOCK / BLOCK_M; + + static constexpr int kBlockM = BLOCK_M; + static constexpr int kBlockN = BLOCK_N; + static constexpr int kHeadDim = D; + + static constexpr int kGmemThreadsPerRow = THREADS_PER_ROW; + static constexpr int kGmemRowsPerThread = 1; + static constexpr int kGmemElemsPerLoad = 8; + + static constexpr int PER_UINT4 = 8; + static constexpr int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; }; -template -__device__ __forceinline__ -int64_t resolve_thread_kv_page_slice_offset( - const int tidx, - const int n_block, - const int page_block_size, - const int* __restrict__ block_table, - const int page_stride, - const int row_stride, - const int partial_block_size = -1 -) { - constexpr int kGmemThreadsPerRow = Traits::kGmemThreadsPerRow; - constexpr int kGmemRowsPerThread = Traits::kGmemRowsPerThread; - constexpr int kGmemElemsPerLoad = Traits::kGmemElemsPerLoad; - constexpr int kBlockN = Traits::kBlockN; - - const int64_t col_offset = (tidx % kGmemThreadsPerRow) * kGmemElemsPerLoad; - int64_t block_row_offset = (tidx / kGmemThreadsPerRow) * kGmemRowsPerThread; - - if (partial_block_size > 0) { - const int final_row_offset = max(partial_block_size - 1, 0); - const int final_thread_row_offset = - ((final_row_offset + kGmemRowsPerThread - 1) / kGmemRowsPerThread) * kGmemRowsPerThread; - block_row_offset = min(block_row_offset, (int64_t)final_thread_row_offset); - } +template +__device__ __forceinline__ int64_t resolve_thread_kv_page_slice_offset( + const int tidx, const int n_block, const int page_block_size, + const int* __restrict__ block_table, const int page_stride, + const int row_stride, const int partial_block_size = -1) { + constexpr int kGmemThreadsPerRow = Traits::kGmemThreadsPerRow; + constexpr int kGmemRowsPerThread = Traits::kGmemRowsPerThread; + constexpr int kGmemElemsPerLoad = Traits::kGmemElemsPerLoad; + constexpr int kBlockN = Traits::kBlockN; + + const int64_t col_offset = (tidx % kGmemThreadsPerRow) * kGmemElemsPerLoad; + int64_t block_row_offset = (tidx / kGmemThreadsPerRow) * kGmemRowsPerThread; + + if (partial_block_size > 0) { + const int final_row_offset = max(partial_block_size - 1, 0); + const int final_thread_row_offset = + ((final_row_offset + kGmemRowsPerThread - 1) / kGmemRowsPerThread) * + kGmemRowsPerThread; + block_row_offset = min(block_row_offset, (int64_t)final_thread_row_offset); + } + + const int64_t global_row_offset = block_row_offset + n_block * kBlockN; + + const int64_t page_offset = global_row_offset % page_block_size; + const int64_t virtual_page_idx = global_row_offset / page_block_size; + + const int64_t physical_offset = + ((int64_t)block_table[virtual_page_idx]) * ((int64_t)page_stride) + + page_offset * ((int64_t)row_stride) + col_offset; + + return physical_offset; +} - const int64_t global_row_offset = block_row_offset + n_block * kBlockN; +template +__device__ __forceinline__ void load_kv_tile_paged( + const __half* __restrict__ kv_cache, const int* __restrict__ block_table, + const int start_token_idx, const int num_tokens, const int n_block, + const int page_block_size, const int page_stride, const int row_stride, + __half* smem_dst, const int smem_stride, const int tid, + const int num_threads, const int partial_block_size = -1) { + constexpr int PER_UINT4 = Traits::PER_UINT4; + constexpr int d_stride_uint4 = Traits::d_stride_uint4; + constexpr int kBlockN = Traits::kBlockN; - const int64_t page_offset = global_row_offset % page_block_size; - const int64_t virtual_page_idx = global_row_offset / page_block_size; + const uint4* kv_vec = reinterpret_cast(kv_cache); + uint4* smem_vec = reinterpret_cast(smem_dst); + const int smem_stride_uint4 = (smem_stride + PER_UINT4 - 1) / PER_UINT4; - const int64_t physical_offset = - ((int64_t)block_table[virtual_page_idx]) * ((int64_t)page_stride) - + page_offset * ((int64_t)row_stride) - + col_offset; +#pragma unroll 2 + for (int idx = tid; idx < (num_tokens * d_stride_uint4); idx += num_threads) { + const int token_offset = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; - return physical_offset; -} + uint4 kv_val = make_uint4(0, 0, 0, 0); -template -__device__ __forceinline__ -void load_kv_tile_paged( - const __half* __restrict__ kv_cache, - const int* __restrict__ block_table, - const int start_token_idx, - const int num_tokens, - const int n_block, - const int page_block_size, - const int page_stride, - const int row_stride, - __half* smem_dst, - const int smem_stride, - const int tid, - const int num_threads, - const int partial_block_size = -1 -) { - constexpr int PER_UINT4 = Traits::PER_UINT4; - constexpr int d_stride_uint4 = Traits::d_stride_uint4; - constexpr int kBlockN = Traits::kBlockN; - - const uint4* kv_vec = reinterpret_cast(kv_cache); - uint4* smem_vec = reinterpret_cast(smem_dst); - const int smem_stride_uint4 = (smem_stride + PER_UINT4 - 1) / PER_UINT4; - - #pragma unroll 2 - for (int idx = tid; idx < (num_tokens * d_stride_uint4); idx += num_threads) { - const int token_offset = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - - uint4 kv_val = make_uint4(0, 0, 0, 0); - - if (token_offset < num_tokens && vec_col < d_stride_uint4) { - - const int global_token_idx = start_token_idx + token_offset; - - const int64_t physical_offset_halfs = resolve_thread_kv_page_slice_offset( - tid, - n_block, - page_block_size, - block_table, - page_stride, - row_stride, - partial_block_size - ); - - const int64_t physical_offset_uint4 = physical_offset_halfs / PER_UINT4; - - kv_val = __ldg(&kv_vec[physical_offset_uint4 + vec_col]); - } - - smem_vec[token_offset * smem_stride_uint4 + vec_col] = kv_val; + if (token_offset < num_tokens && vec_col < d_stride_uint4) { + const int global_token_idx = start_token_idx + token_offset; + + const int64_t physical_offset_halves = + resolve_thread_kv_page_slice_offset( + tid, n_block, page_block_size, block_table, page_stride, + row_stride, partial_block_size); + + const int64_t physical_offset_uint4 = physical_offset_halves / PER_UINT4; + + kv_val = __ldg(&kv_vec[physical_offset_uint4 + vec_col]); } + + smem_vec[token_offset * smem_stride_uint4 + vec_col] = kv_val; + } } diff --git a/flash-attention-v100/kernel/fp8_kv_utils.cuh b/flash-attention-v100/kernel/fp8_kv_utils.cuh index 693b46a93f..dfa54a3eae 100644 --- a/flash-attention-v100/kernel/fp8_kv_utils.cuh +++ b/flash-attention-v100/kernel/fp8_kv_utils.cuh @@ -14,9 +14,7 @@ __device__ __forceinline__ float quiet_nan_f() { return __int_as_float(0x7fffffff); } -__device__ __forceinline__ float inf_f() { - return __int_as_float(0x7f800000); -} +__device__ __forceinline__ float inf_f() { return __int_as_float(0x7f800000); } __device__ __forceinline__ float exp2_int(const int exponent) { return __uint_as_float(static_cast(exponent + 127) << 23); @@ -38,8 +36,7 @@ __device__ __forceinline__ float fp8_e4m3fn_to_float(uint8_t raw) { if (exp == 0x0f && mant == 0x07) { return quiet_nan_f(); } - value = (1.0f + static_cast(mant) * 0.125f) * - exp2_int(exp - 7); + value = (1.0f + static_cast(mant) * 0.125f) * exp2_int(exp - 7); } return sign ? -value : value; } @@ -52,9 +49,8 @@ __device__ __forceinline__ __half fp8_e5m2_to_half(uint8_t raw) { } __device__ __forceinline__ __half2 fp8_e5m2_pair_to_half2(uint16_t raw_pair) { - const uint32_t half2_bits = - (static_cast(raw_pair & 0x00ffu) << 8) | - (static_cast(raw_pair & 0xff00u) << 16); + const uint32_t half2_bits = (static_cast(raw_pair & 0x00ffu) << 8) | + (static_cast(raw_pair & 0xff00u) << 16); union { uint32_t u; __half2 h2; @@ -64,8 +60,7 @@ __device__ __forceinline__ __half2 fp8_e5m2_pair_to_half2(uint16_t raw_pair) { } __device__ __forceinline__ __half2 load_fp8_e5m2_half2_unscaled( - const void* __restrict__ cache, - const int64_t byte_index) { + const void* __restrict__ cache, const int64_t byte_index) { const uint16_t* cache_u16 = reinterpret_cast(cache); return fp8_e5m2_pair_to_half2(cache_u16[byte_index >> 1]); } @@ -74,36 +69,31 @@ __device__ __forceinline__ float fp8_e5m2_to_float(uint8_t raw) { return __half2float(fp8_e5m2_to_half(raw)); } -template +template __device__ __forceinline__ float load_kv_cache_float_unscaled( - const void* __restrict__ cache, - const int64_t index) { + const void* __restrict__ cache, const int64_t index) { if constexpr (KV_DTYPE == KV_CACHE_DTYPE_FP16) { const __half* cache_h = reinterpret_cast(cache); return __half2float(cache_h[index]); } else { const uint8_t* cache_u8 = reinterpret_cast(cache); const uint8_t raw = cache_u8[index]; - const float value = - KV_DTYPE == KV_CACHE_DTYPE_FP8_E4M3 ? fp8_e4m3fn_to_float(raw) - : fp8_e5m2_to_float(raw); + const float value = KV_DTYPE == KV_CACHE_DTYPE_FP8_E4M3 + ? fp8_e4m3fn_to_float(raw) + : fp8_e5m2_to_float(raw); return value; } } -template +template __device__ __forceinline__ float load_kv_cache_float( - const void* __restrict__ cache, - const int64_t index, - const float scale) { + const void* __restrict__ cache, const int64_t index, const float scale) { return load_kv_cache_float_unscaled(cache, index) * scale; } -template +template __device__ __forceinline__ __half load_kv_cache_half( - const void* __restrict__ cache, - const int64_t index, - const float scale) { + const void* __restrict__ cache, const int64_t index, const float scale) { if constexpr (KV_DTYPE == KV_CACHE_DTYPE_FP8_E5M2) { const uint8_t* cache_u8 = reinterpret_cast(cache); return __float2half_rn(__half2float(fp8_e5m2_to_half(cache_u8[index])) * diff --git a/flash-attention-v100/kernel/fused_mha_backward.cu b/flash-attention-v100/kernel/fused_mha_backward.cu index f3e98fed57..1f0ed13072 100644 --- a/flash-attention-v100/kernel/fused_mha_backward.cu +++ b/flash-attention-v100/kernel/fused_mha_backward.cu @@ -14,1176 +14,1235 @@ using namespace nvcuda::wmma; #define WMMA_N 16 #define WMMA_K 16 -#define MAX_THREADS_PER_WARP 32 -#define MAX_THREADS_PER_SM 2048 -#define MAX_THREAD_BLOCK_SIZE 1024 +#define MAX_THREADS_PER_WARP 32 +#define MAX_THREADS_PER_SM 2048 +#define MAX_THREAD_BLOCK_SIZE 1024 #define MAX_THREAD_BLOCK_PER_SM 32 -#define MAX_WARPS_PER_SM 64 -#define MAX_SM_PER_GPU 80 -#define MAX_SMEM_PER_SM 98304 +#define MAX_WARPS_PER_SM 64 +#define MAX_SM_PER_GPU 80 +#define MAX_SMEM_PER_SM 98304 -#define WARP_ALLOC_GROUP 4 +#define WARP_ALLOC_GROUP 4 -#define MAX_REG_PER_UNIT 256 -#define MAX_REG_PER_THREAD 255 -#define MAX_REG_PER_BLOCK 65536 -#define MAX_REG_BUFFER 65536 +#define MAX_REG_PER_UNIT 256 +#define MAX_REG_PER_THREAD 255 +#define MAX_REG_PER_BLOCK 65536 +#define MAX_REG_BUFFER 65536 -#define BLOCK_M_16 16 -#define BLOCK_N_16 256 -#define WARPS_16 16 +#define BLOCK_M_16 16 +#define BLOCK_N_16 256 +#define WARPS_16 16 -#define BLOCK_M_32 32 -#define BLOCK_N_32 128 -#define WARPS_32 16 +#define BLOCK_M_32 32 +#define BLOCK_N_32 128 +#define WARPS_32 16 -#define BLOCK_M_64 64 -#define BLOCK_N_64 80 -#define WARPS_64 16 +#define BLOCK_M_64 64 +#define BLOCK_N_64 80 +#define WARPS_64 16 #define BLOCK_M_128 32 #define BLOCK_N_128 112 -#define WARPS_128 16 +#define WARPS_128 16 #define BLOCK_M_256 32 #define BLOCK_N_256 32 -#define WARPS_256 16 +#define WARPS_256 16 -#define BLOCK_KV_16 32 -#define BLOCK_Q_16 224 +#define BLOCK_KV_16 32 +#define BLOCK_Q_16 224 #define WARPS_DKV_16 14 -#define BLOCK_KV_32 32 -#define BLOCK_Q_32 192 +#define BLOCK_KV_32 32 +#define BLOCK_Q_32 192 #define WARPS_DKV_32 12 -#define BLOCK_KV_64 32 -#define BLOCK_Q_64 128 -#define WARPS_DKV_64 8 +#define BLOCK_KV_64 32 +#define BLOCK_Q_64 128 +#define WARPS_DKV_64 8 -#define BLOCK_KV_128 16 -#define BLOCK_Q_128 144 +#define BLOCK_KV_128 16 +#define BLOCK_Q_128 144 #define WARPS_DKV_128 12 -#define BLOCK_KV_256 16 -#define BLOCK_Q_256 64 +#define BLOCK_KV_256 16 +#define BLOCK_Q_256 64 #define WARPS_DKV_256 16 -template +template struct dQKernelConfig { - static constexpr int BLOCK_M = (D == 16) ? BLOCK_M_16 : (D == 32) ? BLOCK_M_32 : (D == 64) ? BLOCK_M_64 : (D == 128) ? BLOCK_M_128 : BLOCK_M_256; - static constexpr int BLOCK_N = (D == 16) ? BLOCK_N_16 : (D == 32) ? BLOCK_N_32 : (D == 64) ? BLOCK_N_64 : (D == 128) ? BLOCK_N_128 : BLOCK_N_256; - static constexpr int WARPS_PER_BLOCK = (D == 16) ? WARPS_16 : (D == 32) ? WARPS_32 : (D == 64) ? WARPS_64 : (D == 128) ? WARPS_128 : WARPS_256; - - static constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * MAX_THREADS_PER_WARP; - static constexpr int THREADS_PER_ROW = THREADS_PER_BLOCK / BLOCK_M; - static constexpr int PAD = (8 - (D % 32) + 32) % 32; - static constexpr int Q_STRIDE = D + PAD; - static constexpr int KV_STRIDE = D + PAD; - static constexpr int S_STRIDE = BLOCK_N + PAD; - static constexpr int PER_UINT4 = 8; - static constexpr int NUM_UINT4_Q_BLOCK = BLOCK_M * ((D + PER_UINT4 - 1) / PER_UINT4); - static constexpr int NUM_UINT4_KV_BLOCK = BLOCK_N * ((D + PER_UINT4 - 1) / PER_UINT4); - - struct alignas(128) SmemLayout { - union { - alignas(16) __half k [BLOCK_N * KV_STRIDE]; - alignas(16) __half v [BLOCK_N * KV_STRIDE]; - } reuse_kv; - alignas(16) __half dO [BLOCK_M * Q_STRIDE]; - alignas(16) __half q [BLOCK_M * Q_STRIDE]; - alignas(16) float s [BLOCK_M * S_STRIDE]; - union { - alignas(16) float dOV [BLOCK_M * S_STRIDE]; - alignas(16) __half dS [BLOCK_M * S_STRIDE]; - } reuse_sdOVS; - alignas(16) float row_dot[BLOCK_M]; - alignas(16) float lse [BLOCK_M]; - alignas(16) float dQ [BLOCK_M * Q_STRIDE]; - }; - - static constexpr size_t TOTAL_SMEM = ((sizeof(SmemLayout) + 127) & ~size_t(127)); + static constexpr int BLOCK_M = (D == 16) ? BLOCK_M_16 + : (D == 32) ? BLOCK_M_32 + : (D == 64) ? BLOCK_M_64 + : (D == 128) ? BLOCK_M_128 + : BLOCK_M_256; + static constexpr int BLOCK_N = (D == 16) ? BLOCK_N_16 + : (D == 32) ? BLOCK_N_32 + : (D == 64) ? BLOCK_N_64 + : (D == 128) ? BLOCK_N_128 + : BLOCK_N_256; + static constexpr int WARPS_PER_BLOCK = (D == 16) ? WARPS_16 + : (D == 32) ? WARPS_32 + : (D == 64) ? WARPS_64 + : (D == 128) ? WARPS_128 + : WARPS_256; + + static constexpr int THREADS_PER_BLOCK = + WARPS_PER_BLOCK * MAX_THREADS_PER_WARP; + static constexpr int THREADS_PER_ROW = THREADS_PER_BLOCK / BLOCK_M; + static constexpr int PAD = (8 - (D % 32) + 32) % 32; + static constexpr int Q_STRIDE = D + PAD; + static constexpr int KV_STRIDE = D + PAD; + static constexpr int S_STRIDE = BLOCK_N + PAD; + static constexpr int PER_UINT4 = 8; + static constexpr int NUM_UINT4_Q_BLOCK = + BLOCK_M * ((D + PER_UINT4 - 1) / PER_UINT4); + static constexpr int NUM_UINT4_KV_BLOCK = + BLOCK_N * ((D + PER_UINT4 - 1) / PER_UINT4); + + struct alignas(128) SmemLayout { + union { + alignas(16) __half k[BLOCK_N * KV_STRIDE]; + alignas(16) __half v[BLOCK_N * KV_STRIDE]; + } reuse_kv; + alignas(16) __half dO[BLOCK_M * Q_STRIDE]; + alignas(16) __half q[BLOCK_M * Q_STRIDE]; + alignas(16) float s[BLOCK_M * S_STRIDE]; + union { + alignas(16) float dOV[BLOCK_M * S_STRIDE]; + alignas(16) __half dS[BLOCK_M * S_STRIDE]; + } reuse_sdOVS; + alignas(16) float row_dot[BLOCK_M]; + alignas(16) float lse[BLOCK_M]; + alignas(16) float dQ[BLOCK_M * Q_STRIDE]; + }; + + static constexpr size_t TOTAL_SMEM = + ((sizeof(SmemLayout) + 127) & ~size_t(127)); }; -template +template __device__ __forceinline__ void init_smem(char* smem_raw) { - constexpr int N_U4 = Config::TOTAL_SMEM / 16; - const int lane_id = threadIdx.x & 31; - - uint32_t addr = static_cast(__cvta_generic_to_shared(smem_raw)); - #pragma unroll 1 - for (int i = lane_id; i < N_U4; i += 32) { - asm volatile("st.shared.v4.u32 [%0], {%1,%1,%1,%1};" - :: "r"(addr + (i << 4)), "r"(0) : "memory"); - } - __syncwarp(); + constexpr int N_U4 = Config::TOTAL_SMEM / 16; + const int lane_id = threadIdx.x & 31; + + uint32_t addr = static_cast(__cvta_generic_to_shared(smem_raw)); +#pragma unroll 1 + for (int i = lane_id; i < N_U4; i += 32) { + asm volatile("st.shared.v4.u32 [%0], {%1,%1,%1,%1};" ::"r"(addr + (i << 4)), + "r"(0) + : "memory"); + } + __syncwarp(); } -template +template struct dKVKernelConfig { - static constexpr int BLOCK_M = (D == 16) ? BLOCK_KV_16 : (D == 32) ? BLOCK_KV_32 : (D == 64) ? BLOCK_KV_64 : (D == 128) ? BLOCK_KV_128 : BLOCK_KV_256; - static constexpr int BLOCK_N = (D == 16) ? BLOCK_Q_16 : (D == 32) ? BLOCK_Q_32 : (D == 64) ? BLOCK_Q_64 : (D == 128) ? BLOCK_Q_128 : BLOCK_Q_256; - static constexpr int WARPS_PER_BLOCK = (D == 16) ? WARPS_DKV_16 : (D == 32) ? WARPS_DKV_32 : (D == 64) ? WARPS_DKV_64 : (D == 128) ? WARPS_DKV_128 : WARPS_DKV_256; - - static constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * MAX_THREADS_PER_WARP; - static constexpr int THREADS_PER_ROW = THREADS_PER_BLOCK / BLOCK_N; - static constexpr int PAD = 8; - static constexpr int Q_STRIDE = D + PAD; - static constexpr int KV_STRIDE = D + PAD; - static constexpr int S_STRIDE = BLOCK_M + PAD; - static constexpr int PER_UINT4 = 8; - static constexpr int NUM_UINT4_Q_BLOCK = BLOCK_N * ((D + PER_UINT4 - 1) / PER_UINT4); - static constexpr int NUM_UINT4_KV_BLOCK = BLOCK_M * ((D + PER_UINT4 - 1) / PER_UINT4); - - struct alignas(128) SmemLayout { - alignas(16) __half k [BLOCK_M * KV_STRIDE]; - alignas(16) __half v [BLOCK_M * KV_STRIDE]; - union { - alignas(16) __half dO [BLOCK_N * Q_STRIDE]; - alignas(16) __half q [BLOCK_N * Q_STRIDE]; - } reuse_qdO; - union { - alignas(16) float s [BLOCK_N * S_STRIDE]; - alignas(16) __half p [BLOCK_N * BLOCK_M]; - } reuse_sp; - union { - alignas(16) float dOV [BLOCK_N * S_STRIDE]; - alignas(16) __half dS [BLOCK_N * BLOCK_M]; - } reuse_dOVS; - alignas(16) float row_dot[BLOCK_N]; - alignas(16) float lse [BLOCK_N]; - alignas(16) float dK [BLOCK_M * KV_STRIDE]; - alignas(16) float dV [BLOCK_M * KV_STRIDE]; - }; - - static constexpr size_t TOTAL_SMEM = ((sizeof(SmemLayout) + 127) & ~size_t(127)); + static constexpr int BLOCK_M = (D == 16) ? BLOCK_KV_16 + : (D == 32) ? BLOCK_KV_32 + : (D == 64) ? BLOCK_KV_64 + : (D == 128) ? BLOCK_KV_128 + : BLOCK_KV_256; + static constexpr int BLOCK_N = (D == 16) ? BLOCK_Q_16 + : (D == 32) ? BLOCK_Q_32 + : (D == 64) ? BLOCK_Q_64 + : (D == 128) ? BLOCK_Q_128 + : BLOCK_Q_256; + static constexpr int WARPS_PER_BLOCK = (D == 16) ? WARPS_DKV_16 + : (D == 32) ? WARPS_DKV_32 + : (D == 64) ? WARPS_DKV_64 + : (D == 128) ? WARPS_DKV_128 + : WARPS_DKV_256; + + static constexpr int THREADS_PER_BLOCK = + WARPS_PER_BLOCK * MAX_THREADS_PER_WARP; + static constexpr int THREADS_PER_ROW = THREADS_PER_BLOCK / BLOCK_N; + static constexpr int PAD = 8; + static constexpr int Q_STRIDE = D + PAD; + static constexpr int KV_STRIDE = D + PAD; + static constexpr int S_STRIDE = BLOCK_M + PAD; + static constexpr int PER_UINT4 = 8; + static constexpr int NUM_UINT4_Q_BLOCK = + BLOCK_N * ((D + PER_UINT4 - 1) / PER_UINT4); + static constexpr int NUM_UINT4_KV_BLOCK = + BLOCK_M * ((D + PER_UINT4 - 1) / PER_UINT4); + + struct alignas(128) SmemLayout { + alignas(16) __half k[BLOCK_M * KV_STRIDE]; + alignas(16) __half v[BLOCK_M * KV_STRIDE]; + union { + alignas(16) __half dO[BLOCK_N * Q_STRIDE]; + alignas(16) __half q[BLOCK_N * Q_STRIDE]; + } reuse_qdO; + union { + alignas(16) float s[BLOCK_N * S_STRIDE]; + alignas(16) __half p[BLOCK_N * BLOCK_M]; + } reuse_sp; + union { + alignas(16) float dOV[BLOCK_N * S_STRIDE]; + alignas(16) __half dS[BLOCK_N * BLOCK_M]; + } reuse_dOVS; + alignas(16) float row_dot[BLOCK_N]; + alignas(16) float lse[BLOCK_N]; + alignas(16) float dK[BLOCK_M * KV_STRIDE]; + alignas(16) float dV[BLOCK_M * KV_STRIDE]; + }; + + static constexpr size_t TOTAL_SMEM = + ((sizeof(SmemLayout) + 127) & ~size_t(127)); }; -template +template __global__ void __launch_bounds__(dQKernelConfig::THREADS_PER_BLOCK, 2) -flash_attention_backward_dq_kernel( - const __half* __restrict__ Q, - const __half* __restrict__ K, - const __half* __restrict__ V, - const __half* __restrict__ O, - const __half* __restrict__ dO, - const float* __restrict__ softmax_lse, - __half* __restrict__ dQ, - const int B, - const int H, - const int M, - const int N, - const float softmax_scale -) { - using Config = dQKernelConfig; - - constexpr int BLOCK_M = Config::BLOCK_M; - constexpr int BLOCK_N = Config::BLOCK_N; - constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; - constexpr int THREADS_PER_ROW = Config::THREADS_PER_ROW; - constexpr int WARPS_PER_BLOCK = Config::WARPS_PER_BLOCK; - constexpr int Q_STRIDE = Config::Q_STRIDE; - constexpr int KV_STRIDE = Config::KV_STRIDE; - constexpr int S_STRIDE = Config::S_STRIDE; - constexpr int PER_UINT4 = Config::PER_UINT4; - constexpr int NUM_UINT4_Q_BLOCK = Config::NUM_UINT4_Q_BLOCK; - constexpr int NUM_UINT4_KV_BLOCK = Config::NUM_UINT4_KV_BLOCK; - - const float NEG_INF = -1e30f; - - const int batch_head_id = blockIdx.z; - if (batch_head_id >= B * H) return; - - const int block_m = blockIdx.x; - const int start_row = block_m * BLOCK_M; - if (start_row >= M) return; - - int num_n_tiles = (N + BLOCK_N - 1) / BLOCK_N; - const int valid_q_rows = min(BLOCK_M, M - start_row); - - if constexpr (IS_CAUSAL) { - const int max_key_pos = start_row + valid_q_rows - 1; - if (max_key_pos < 0) { - num_n_tiles = 0; - } else { - num_n_tiles = min(num_n_tiles, (max_key_pos + BLOCK_N) / BLOCK_N); - } + flash_attention_backward_dq_kernel( + const __half* __restrict__ Q, const __half* __restrict__ K, + const __half* __restrict__ V, const __half* __restrict__ O, + const __half* __restrict__ dO, const float* __restrict__ softmax_lse, + __half* __restrict__ dQ, const int B, const int H, const int M, + const int N, const float softmax_scale) { + using Config = dQKernelConfig; + + constexpr int BLOCK_M = Config::BLOCK_M; + constexpr int BLOCK_N = Config::BLOCK_N; + constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; + constexpr int THREADS_PER_ROW = Config::THREADS_PER_ROW; + constexpr int WARPS_PER_BLOCK = Config::WARPS_PER_BLOCK; + constexpr int Q_STRIDE = Config::Q_STRIDE; + constexpr int KV_STRIDE = Config::KV_STRIDE; + constexpr int S_STRIDE = Config::S_STRIDE; + constexpr int PER_UINT4 = Config::PER_UINT4; + constexpr int NUM_UINT4_Q_BLOCK = Config::NUM_UINT4_Q_BLOCK; + constexpr int NUM_UINT4_KV_BLOCK = Config::NUM_UINT4_KV_BLOCK; + + const float NEG_INF = -1e30f; + + const int batch_head_id = blockIdx.z; + if (batch_head_id >= B * H) return; + + const int block_m = blockIdx.x; + const int start_row = block_m * BLOCK_M; + if (start_row >= M) return; + + int num_n_tiles = (N + BLOCK_N - 1) / BLOCK_N; + const int valid_q_rows = min(BLOCK_M, M - start_row); + + if constexpr (IS_CAUSAL) { + const int max_key_pos = start_row + valid_q_rows - 1; + if (max_key_pos < 0) { + num_n_tiles = 0; + } else { + num_n_tiles = min(num_n_tiles, (max_key_pos + BLOCK_N) / BLOCK_N); + } + } + + const int tid = threadIdx.x; + const int warp_id = tid / MAX_THREADS_PER_WARP; + const int lane_id = tid % MAX_THREADS_PER_WARP; + + const __half* q_ptr = Q + (size_t)batch_head_id * M * D + start_row * D; + const __half* k_ptr = K + (size_t)batch_head_id * N * D; + const __half* v_ptr = V + (size_t)batch_head_id * N * D; + const __half* o_ptr = O + (size_t)batch_head_id * M * D + start_row * D; + const __half* dO_ptr = dO + (size_t)batch_head_id * M * D + start_row * D; + __half* dQ_ptr = dQ + (size_t)batch_head_id * M * D + start_row * D; + const float* lse_ptr = softmax_lse + (size_t)batch_head_id * M + start_row; + + extern __shared__ char smem_raw[]; + init_smem(smem_raw); + auto& smem = *reinterpret_cast(smem_raw); + + __half* sK = smem.reuse_kv.k; + __half* sV = smem.reuse_kv.v; + __half* sdO = smem.dO; + __half* sQ = smem.q; + float* sS = smem.s; + float* sdOV = smem.reuse_sdOVS.dOV; + __half* sdS = smem.reuse_sdOVS.dS; + float* sRowDot = smem.row_dot; + float* sLse = smem.lse; + float* sdQ = smem.dQ; + + const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; + const int q_stride_uint4 = (Q_STRIDE + PER_UINT4 - 1) / PER_UINT4; + const int kv_stride_uint4 = (KV_STRIDE + PER_UINT4 - 1) / PER_UINT4; + + const uint4* q_vec = reinterpret_cast(q_ptr); + const uint4* do_vec = reinterpret_cast(dO_ptr); + uint4* sQ_vec = reinterpret_cast(sQ); + uint4* sdO_vec = reinterpret_cast(sdO); + +#pragma unroll 2 + for (int idx = tid; idx < NUM_UINT4_Q_BLOCK; idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + + uint4 q_val = make_uint4(0, 0, 0, 0); + uint4 do_val = make_uint4(0, 0, 0, 0); + + if (row < valid_q_rows && vec_col < d_stride_uint4) { + q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); + do_val = __ldg(&do_vec[row * d_stride_uint4 + vec_col]); + } + sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + sdO_vec[row * q_stride_uint4 + vec_col] = do_val; + } + __syncthreads(); + + if (tid < valid_q_rows * THREADS_PER_ROW) { + const int row = tid / THREADS_PER_ROW; + const int thread_in_row = tid % THREADS_PER_ROW; + const int fp16_x4_per_row = D / 4; + const int work_per_thread = + (fp16_x4_per_row + THREADS_PER_ROW - 1) / THREADS_PER_ROW; + const unsigned mask = + (valid_q_rows == BLOCK_M) ? 0xFFFFFFFFU : __activemask(); + + float thread_dot = 0.0f; + +#pragma unroll + for (int j = 0; j < work_per_thread; ++j) { + const int chunk_idx = thread_in_row + j * THREADS_PER_ROW; + if (chunk_idx >= fp16_x4_per_row) break; + const int col = chunk_idx * 4; + + const __half* o_addr = o_ptr + row * D + col; + ushort o_h0, o_h1, o_h2, o_h3; + asm volatile("ld.global.v4.u16 {%0, %1, %2, %3}, [%4];" + : "=h"(o_h0), "=h"(o_h1), "=h"(o_h2), "=h"(o_h3) + : "l"(o_addr) + : "memory"); + + const __half* dO_addr = sdO + row * Q_STRIDE + col; + ushort d_h0, d_h1, d_h2, d_h3; + const uint32_t ptr_dO = static_cast( + reinterpret_cast(__cvta_generic_to_shared(dO_addr))); + asm volatile("ld.shared.v4.u16 {%0, %1, %2, %3}, [%4];" + : "=h"(d_h0), "=h"(d_h1), "=h"(d_h2), "=h"(d_h3) + : "r"(ptr_dO) + : "memory"); + + const float fo_0 = __half2float(__ushort_as_half(o_h0)); + const float fo_1 = __half2float(__ushort_as_half(o_h1)); + const float fo_2 = __half2float(__ushort_as_half(o_h2)); + const float fo_3 = __half2float(__ushort_as_half(o_h3)); + + const float fd_0 = __half2float(__ushort_as_half(d_h0)); + const float fd_1 = __half2float(__ushort_as_half(d_h1)); + const float fd_2 = __half2float(__ushort_as_half(d_h2)); + const float fd_3 = __half2float(__ushort_as_half(d_h3)); + + thread_dot = __fmaf_rn(fo_0, fd_0, thread_dot); + thread_dot = __fmaf_rn(fo_1, fd_1, thread_dot); + thread_dot = __fmaf_rn(fo_2, fd_2, thread_dot); + thread_dot = __fmaf_rn(fo_3, fd_3, thread_dot); } - const int tid = threadIdx.x; - const int warp_id = tid / MAX_THREADS_PER_WARP; - const int lane_id = tid % MAX_THREADS_PER_WARP; - - const __half* q_ptr = Q + (size_t)batch_head_id * M * D + start_row * D; - const __half* k_ptr = K + (size_t)batch_head_id * N * D; - const __half* v_ptr = V + (size_t)batch_head_id * N * D; - const __half* o_ptr = O + (size_t)batch_head_id * M * D + start_row * D; - const __half* dO_ptr = dO + (size_t)batch_head_id * M * D + start_row * D; - __half* dQ_ptr = dQ + (size_t)batch_head_id * M * D + start_row * D; - const float* lse_ptr = softmax_lse + (size_t)batch_head_id * M + start_row; - - extern __shared__ char smem_raw[]; - init_smem(smem_raw); - auto& smem = *reinterpret_cast(smem_raw); - - __half* sK = smem.reuse_kv.k; - __half* sV = smem.reuse_kv.v; - __half* sdO = smem.dO; - __half* sQ = smem.q; - float* sS = smem.s; - float* sdOV = smem.reuse_sdOVS.dOV; - __half* sdS = smem.reuse_sdOVS.dS; - float* sRowDot = smem.row_dot; - float* sLse = smem.lse; - float* sdQ = smem.dQ; - - const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; - const int q_stride_uint4 = (Q_STRIDE + PER_UINT4 - 1) / PER_UINT4; - const int kv_stride_uint4 = (KV_STRIDE + PER_UINT4 - 1) / PER_UINT4; - - const uint4* q_vec = reinterpret_cast(q_ptr); - const uint4* do_vec = reinterpret_cast(dO_ptr); - uint4* sQ_vec = reinterpret_cast(sQ); - uint4* sdO_vec = reinterpret_cast(sdO); - - #pragma unroll 2 - for (int idx = tid; idx < NUM_UINT4_Q_BLOCK; idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - - uint4 q_val = make_uint4(0, 0, 0, 0); - uint4 do_val = make_uint4(0, 0, 0, 0); +#pragma unroll + for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) + thread_dot += __shfl_down_sync(mask, thread_dot, o, THREADS_PER_ROW); - if (row < valid_q_rows && vec_col < d_stride_uint4) { - q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); - do_val = __ldg(&do_vec[row * d_stride_uint4 + vec_col]); - } - sQ_vec[row * q_stride_uint4 + vec_col] = q_val; - sdO_vec[row * q_stride_uint4 + vec_col] = do_val; + if (thread_in_row == 0) { + sRowDot[row] = thread_dot; } - __syncthreads(); + } - if (tid < valid_q_rows * THREADS_PER_ROW) { - const int row = tid / THREADS_PER_ROW; - const int thread_in_row = tid % THREADS_PER_ROW; - const int fp16_x4_per_row = D / 4; - const int work_per_thread = (fp16_x4_per_row + THREADS_PER_ROW - 1) / THREADS_PER_ROW; - const unsigned mask = (valid_q_rows == BLOCK_M) ? 0xFFFFFFFFU : __activemask(); - - float thread_dot = 0.0f; - - #pragma unroll - for (int j = 0; j < work_per_thread; ++j) { - const int chunk_idx = thread_in_row + j * THREADS_PER_ROW; - if (chunk_idx >= fp16_x4_per_row) break; - const int col = chunk_idx * 4; - - const __half* o_addr = o_ptr + row * D + col; - ushort o_h0, o_h1, o_h2, o_h3; - asm volatile( - "ld.global.v4.u16 {%0, %1, %2, %3}, [%4];" - : "=h"(o_h0), "=h"(o_h1), "=h"(o_h2), "=h"(o_h3) - : "l"(o_addr) - : "memory" - ); - - const __half* dO_addr = sdO + row * Q_STRIDE + col; - ushort d_h0, d_h1, d_h2, d_h3; - const uint32_t ptr_dO = static_cast(reinterpret_cast(__cvta_generic_to_shared(dO_addr))); - asm volatile( - "ld.shared.v4.u16 {%0, %1, %2, %3}, [%4];" - : "=h"(d_h0), "=h"(d_h1), "=h"(d_h2), "=h"(d_h3) - : "r"(ptr_dO) - : "memory" - ); - - const float fo_0 = __half2float(__ushort_as_half(o_h0)); - const float fo_1 = __half2float(__ushort_as_half(o_h1)); - const float fo_2 = __half2float(__ushort_as_half(o_h2)); - const float fo_3 = __half2float(__ushort_as_half(o_h3)); - - const float fd_0 = __half2float(__ushort_as_half(d_h0)); - const float fd_1 = __half2float(__ushort_as_half(d_h1)); - const float fd_2 = __half2float(__ushort_as_half(d_h2)); - const float fd_3 = __half2float(__ushort_as_half(d_h3)); - - thread_dot = __fmaf_rn(fo_0, fd_0, thread_dot); - thread_dot = __fmaf_rn(fo_1, fd_1, thread_dot); - thread_dot = __fmaf_rn(fo_2, fd_2, thread_dot); - thread_dot = __fmaf_rn(fo_3, fd_3, thread_dot); - } + if (tid < valid_q_rows) { + sLse[tid] = lse_ptr[tid]; + } + __syncthreads(); - #pragma unroll - for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) - thread_dot += __shfl_down_sync(mask, thread_dot, o, THREADS_PER_ROW); + for (int block_n = 0; block_n < num_n_tiles; ++block_n) { + const int start_col = block_n * BLOCK_N; + if (start_col >= N) break; + const int valid_k_rows = min(BLOCK_N, N - start_col); - if (thread_in_row == 0) { - sRowDot[row] = thread_dot; - } + if constexpr (IS_CAUSAL) { + if (start_col >= start_row + valid_q_rows) { + continue; + } } - if (tid < valid_q_rows) { sLse[tid] = lse_ptr[tid]; } + const uint4* v_vec = reinterpret_cast(v_ptr + start_col * D); + uint4* sV_vec = reinterpret_cast(sV); + +#pragma unroll 2 + for (int idx = tid; idx < NUM_UINT4_KV_BLOCK; idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + uint4 v_val = make_uint4(0, 0, 0, 0); + if (row < valid_k_rows && vec_col < d_stride_uint4) { + v_val = __ldg(&v_vec[row * d_stride_uint4 + vec_col]); + } + sV_vec[row * kv_stride_uint4 + vec_col] = v_val; + } __syncthreads(); - for (int block_n = 0; block_n < num_n_tiles; ++block_n) { - const int start_col = block_n * BLOCK_N; - if (start_col >= N) break; - const int valid_k_rows = min(BLOCK_N, N - start_col); + const int num_tiles_m_dov = (BLOCK_M + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_dov = (BLOCK_N + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_dov = (D + WMMA_K - 1) / WMMA_K; + const int total_tiles_dov = num_tiles_m_dov * num_tiles_n_dov; + const int tiles_per_warp_dov = + (total_tiles_dov + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - if constexpr (IS_CAUSAL) { - if (start_col >= start_row + valid_q_rows) { continue; } - } + for (int tile_idx = 0; tile_idx < tiles_per_warp_dov; ++tile_idx) { + const int global_tile_idx = warp_id * tiles_per_warp_dov + tile_idx; - const uint4* v_vec = reinterpret_cast(v_ptr + start_col * D); - uint4* sV_vec = reinterpret_cast(sV); - - #pragma unroll 2 - for (int idx = tid; idx < NUM_UINT4_KV_BLOCK; idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - uint4 v_val = make_uint4(0, 0, 0, 0); - if (row < valid_k_rows && vec_col < d_stride_uint4) { - v_val = __ldg(&v_vec[row * d_stride_uint4 + vec_col]); - } - sV_vec[row * kv_stride_uint4 + vec_col] = v_val; - } - __syncthreads(); + if (global_tile_idx >= total_tiles_dov) break; - const int num_tiles_m_dov = (BLOCK_M + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_dov = (BLOCK_N + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_dov = (D + WMMA_K - 1) / WMMA_K; - const int total_tiles_dov = num_tiles_m_dov * num_tiles_n_dov; - const int tiles_per_warp_dov = (total_tiles_dov + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + const int tile_m_idx = global_tile_idx / num_tiles_n_dov; + const int tile_n_idx = global_tile_idx % num_tiles_n_dov; - for (int tile_idx = 0; tile_idx < tiles_per_warp_dov; ++tile_idx) { - const int global_tile_idx = warp_id * tiles_per_warp_dov + tile_idx; + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; - if (global_tile_idx >= total_tiles_dov) break; + if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; - const int tile_m_idx = global_tile_idx / num_tiles_n_dov; - const int tile_n_idx = global_tile_idx % num_tiles_n_dov; + fragment a_frag; + fragment b_frag; + fragment acc_frag; - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; + fill_fragment(acc_frag, 0.0f); - if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; +#pragma unroll + for (int k_tile = 0; k_tile < num_tiles_k_dov; ++k_tile) { + const int k_offset = k_tile * WMMA_K; + if (k_offset >= D) break; + load_matrix_sync(a_frag, sdO + tile_m * Q_STRIDE + k_offset, Q_STRIDE); + load_matrix_sync(b_frag, sV + tile_n * KV_STRIDE + k_offset, KV_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } - fragment a_frag; - fragment b_frag; - fragment acc_frag; + store_matrix_sync(sdOV + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, + mem_row_major); + } + __syncthreads(); - fill_fragment(acc_frag, 0.0f); + const uint4* k_vec = reinterpret_cast(k_ptr + start_col * D); + uint4* sK_vec = reinterpret_cast(sK); - #pragma unroll - for (int k_tile = 0; k_tile < num_tiles_k_dov; ++k_tile) { - const int k_offset = k_tile * WMMA_K; - if (k_offset >= D) break; - load_matrix_sync(a_frag, sdO + tile_m * Q_STRIDE + k_offset, Q_STRIDE); - load_matrix_sync(b_frag, sV + tile_n * KV_STRIDE + k_offset, KV_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } +#pragma unroll 2 + for (int idx = tid; idx < NUM_UINT4_KV_BLOCK; idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + uint4 k_val = make_uint4(0, 0, 0, 0); + if (row < valid_k_rows && vec_col < d_stride_uint4) { + k_val = __ldg(&k_vec[row * d_stride_uint4 + vec_col]); + } + sK_vec[row * kv_stride_uint4 + vec_col] = k_val; + } + __syncthreads(); - store_matrix_sync(sdOV + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, mem_row_major); - } - __syncthreads(); - - const uint4* k_vec = reinterpret_cast(k_ptr + start_col * D); - uint4* sK_vec = reinterpret_cast(sK); - - #pragma unroll 2 - for (int idx = tid; idx < NUM_UINT4_KV_BLOCK; idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - uint4 k_val = make_uint4(0, 0, 0, 0); - if (row < valid_k_rows && vec_col < d_stride_uint4) { - k_val = __ldg(&k_vec[row * d_stride_uint4 + vec_col]); - } - sK_vec[row * kv_stride_uint4 + vec_col] = k_val; + const int num_tiles_m_qk = (BLOCK_M + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_qk = (BLOCK_N + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_qk = (D + WMMA_K - 1) / WMMA_K; + const int total_tiles_qk = num_tiles_m_qk * num_tiles_n_qk; + const int tiles_per_warp_qk = + (total_tiles_qk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + const unsigned row_causal = (lane_id & 0b1) + ((lane_id >> 2) & 0b1) * 8 + + ((lane_id >> 4) & 0b1) * 4; + const unsigned col_causal = + ((lane_id >> 1) & 0b1) * 2 + ((lane_id >> 3) & 0b1) * 8; + + for (int tile_idx = 0; tile_idx < tiles_per_warp_qk; ++tile_idx) { + const int global_tile_idx = warp_id * tiles_per_warp_qk + tile_idx; + + if (global_tile_idx >= total_tiles_qk) break; + + const int tile_m_idx = global_tile_idx / num_tiles_n_qk; + const int tile_n_idx = global_tile_idx % num_tiles_n_qk; + + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; + + if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + + fill_fragment(acc_frag, 0.0f); + +#pragma unroll + for (int k_tile = 0; k_tile < num_tiles_k_qk; ++k_tile) { + const int k_offset = k_tile * WMMA_K; + if (k_offset >= D) break; + load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, Q_STRIDE); + load_matrix_sync(b_frag, sK + tile_n * KV_STRIDE + k_offset, KV_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + if constexpr (IS_CAUSAL) { +#pragma unroll + for (int i = 0; i < acc_frag.num_elements; ++i) { + const unsigned col = col_causal + (i & 0b1) + ((i >> 2) & 0b1) * 4; + const unsigned row = row_causal + ((i >> 1) & 0b1) * 2; + + const int global_m = start_row + tile_m + row; + const int global_n = start_col + tile_n + col; + + const bool is_valid = (global_m < start_row + valid_q_rows) && + (global_n < start_col + valid_k_rows); + + acc_frag.x[i] = + is_valid ? ((global_n > global_m) ? NEG_INF + : acc_frag.x[i] * softmax_scale) + : NEG_INF; } - __syncthreads(); - - const int num_tiles_m_qk = (BLOCK_M + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_qk = (BLOCK_N + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_qk = (D + WMMA_K - 1) / WMMA_K; - const int total_tiles_qk = num_tiles_m_qk * num_tiles_n_qk; - const int tiles_per_warp_qk = (total_tiles_qk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - const unsigned row_causal = (lane_id & 0b1) + ((lane_id >> 2) & 0b1) * 8 + ((lane_id >> 4) & 0b1) * 4; - const unsigned col_causal = ((lane_id >> 1) & 0b1) * 2 + ((lane_id >> 3) & 0b1) * 8; - - for (int tile_idx = 0; tile_idx < tiles_per_warp_qk; ++tile_idx) { - const int global_tile_idx = warp_id * tiles_per_warp_qk + tile_idx; - - if (global_tile_idx >= total_tiles_qk) break; - - const int tile_m_idx = global_tile_idx / num_tiles_n_qk; - const int tile_n_idx = global_tile_idx % num_tiles_n_qk; - - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; - - if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - - fill_fragment(acc_frag, 0.0f); - - #pragma unroll - for (int k_tile = 0; k_tile < num_tiles_k_qk; ++k_tile) { - const int k_offset = k_tile * WMMA_K; - if (k_offset >= D) break; - load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, Q_STRIDE); - load_matrix_sync(b_frag, sK + tile_n * KV_STRIDE + k_offset, KV_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - - if constexpr (IS_CAUSAL) { - #pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - const unsigned col = col_causal + (i & 0b1) + ((i >> 2) & 0b1) * 4; - const unsigned row = row_causal + ((i >> 1) & 0b1) * 2; - - const int global_m = start_row + tile_m + row; - const int global_n = start_col + tile_n + col; - - const bool is_valid = (global_m < start_row + valid_q_rows) && - (global_n < start_col + valid_k_rows); - - acc_frag.x[i] = is_valid - ? ((global_n > global_m) ? NEG_INF : acc_frag.x[i] * softmax_scale) - : NEG_INF; - } - } else { - #pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - acc_frag.x[i] *= softmax_scale; - } - } - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, mem_row_major); + } else { +#pragma unroll + for (int i = 0; i < acc_frag.num_elements; ++i) { + acc_frag.x[i] *= softmax_scale; } - __syncthreads(); - - if (tid < valid_q_rows * THREADS_PER_ROW) { - const int row = tid / THREADS_PER_ROW; - const int thread_in_row = tid % THREADS_PER_ROW; - - float* sS_row = sS + row * S_STRIDE; - float* sdOV_row = sdOV + row * S_STRIDE; - __half* sdS_row = sdS + row * S_STRIDE; - - const float lse_val = sLse[row]; - const float row_dot_val = sRowDot[row]; - - const int vec8_cols = valid_k_rows / 8; - const int vec8_per_thread = (vec8_cols + THREADS_PER_ROW - 1) / THREADS_PER_ROW; - const int tail_start = vec8_cols * 8; - - float4* sS_vec4 = reinterpret_cast(sS_row); - float4* sdOV_vec4 = reinterpret_cast(sdOV_row); - uint4* sdS_vec_u4 = reinterpret_cast(sdS_row); - - uint4 buf[8]; int cnt = 0; - - #pragma unroll - for (int j = 0; j < vec8_per_thread; ++j) { - const int v8 = thread_in_row + j * THREADS_PER_ROW; - if (v8 >= vec8_cols) break; - - float4 s0 = sS_vec4[v8 * 2], s1 = sS_vec4[v8 * 2 + 1]; - float4 d0 = sdOV_vec4[v8 * 2], d1 = sdOV_vec4[v8 * 2 + 1]; - - #define COMP(i, sf, df) \ - float sh##i = (sf) - lse_val; \ - float p##i = (sh##i < -80.0f) ? 0.0f : __expf(sh##i); \ - float ds##i = p##i * softmax_scale * ((df) - row_dot_val); - - COMP(0, s0.x, d0.x) COMP(1, s0.y, d0.y) COMP(2, s0.z, d0.z) COMP(3, s0.w, d0.w) - COMP(4, s1.x, d1.x) COMP(5, s1.y, d1.y) COMP(6, s1.z, d1.z) COMP(7, s1.w, d1.w) - #undef COMP - - uint4 res; - asm volatile( - "{ mov.b32 %0, {%4,%5}; mov.b32 %1, {%6,%7}; mov.b32 %2, {%8,%9}; mov.b32 %3, {%10,%11}; }\n" - : "=r"(res.x), "=r"(res.y), "=r"(res.z), "=r"(res.w) - : "h"(__half_as_ushort(__float2half_rn(ds0))), - "h"(__half_as_ushort(__float2half_rn(ds1))), - "h"(__half_as_ushort(__float2half_rn(ds2))), - "h"(__half_as_ushort(__float2half_rn(ds3))), - "h"(__half_as_ushort(__float2half_rn(ds4))), - "h"(__half_as_ushort(__float2half_rn(ds5))), - "h"(__half_as_ushort(__float2half_rn(ds6))), - "h"(__half_as_ushort(__float2half_rn(ds7))) - ); - buf[cnt++] = res; - } - - #pragma unroll - for (int c = tail_start + thread_in_row; c < BLOCK_N; c += THREADS_PER_ROW) { - float s = (c < valid_k_rows) ? sS_row[c] : NEG_INF; - float dov = (c < valid_k_rows) ? sdOV_row[c] : 0.0f; - float p = (s - lse_val < -80.0f) ? 0.0f : __expf(s - lse_val); - float ds = p * softmax_scale * ((c < valid_k_rows) ? (dov - row_dot_val) : 0.0f); - sdS_row[c] = __float2half_rn(ds); - } - - #pragma unroll - for (int i = 0; i < cnt; ++i) { - const int v8 = thread_in_row + i * THREADS_PER_ROW; - if (v8 < vec8_cols) { - sdS_vec_u4[v8] = buf[i]; - } - } - } - __syncthreads(); - - const int num_tiles_m_dq = (BLOCK_M + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_dq = (D + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_dq = (BLOCK_N + WMMA_K - 1) / WMMA_K; - const int total_tiles_dq = num_tiles_m_dq * num_tiles_n_dq; - const int tiles_per_warp_dq = (total_tiles_dq + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - - for (int tile_local = 0; tile_local < tiles_per_warp_dq; ++tile_local) { - const int global_tile_idx = warp_id * tiles_per_warp_dq + tile_local; - if (global_tile_idx >= total_tiles_dq) break; - - const int tile_m_idx = global_tile_idx / num_tiles_n_dq; - const int tile_n_idx = global_tile_idx % num_tiles_n_dq; - - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; - - if (tile_m >= valid_q_rows || tile_n >= D) continue; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - fill_fragment(acc_frag, 0.0f); - - #pragma unroll - for (int k_tile = 0; k_tile < num_tiles_k_dq; ++k_tile) { - const int k_offset = k_tile * WMMA_K; - if (k_offset >= valid_k_rows) break; - - load_matrix_sync(a_frag, sdS + tile_m * S_STRIDE + k_offset, S_STRIDE); - load_matrix_sync(b_frag, sK + k_offset * KV_STRIDE + tile_n, KV_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - - fragment curr_frag; - load_matrix_sync(curr_frag, sdQ + tile_m * Q_STRIDE + tile_n, Q_STRIDE, mem_row_major); + } + store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, + mem_row_major); + } + __syncthreads(); - #pragma unroll - for (int i = 0; i < curr_frag.num_elements; ++i) { - curr_frag.x[i] += acc_frag.x[i]; - } - store_matrix_sync(sdQ + tile_m * Q_STRIDE + tile_n, curr_frag, Q_STRIDE, mem_row_major); + if (tid < valid_q_rows * THREADS_PER_ROW) { + const int row = tid / THREADS_PER_ROW; + const int thread_in_row = tid % THREADS_PER_ROW; + + float* sS_row = sS + row * S_STRIDE; + float* sdOV_row = sdOV + row * S_STRIDE; + __half* sdS_row = sdS + row * S_STRIDE; + + const float lse_val = sLse[row]; + const float row_dot_val = sRowDot[row]; + + const int vec8_cols = valid_k_rows / 8; + const int vec8_per_thread = + (vec8_cols + THREADS_PER_ROW - 1) / THREADS_PER_ROW; + const int tail_start = vec8_cols * 8; + + float4* sS_vec4 = reinterpret_cast(sS_row); + float4* sdOV_vec4 = reinterpret_cast(sdOV_row); + uint4* sdS_vec_u4 = reinterpret_cast(sdS_row); + + uint4 buf[8]; + int cnt = 0; + +#pragma unroll + for (int j = 0; j < vec8_per_thread; ++j) { + const int v8 = thread_in_row + j * THREADS_PER_ROW; + if (v8 >= vec8_cols) break; + + float4 s0 = sS_vec4[v8 * 2], s1 = sS_vec4[v8 * 2 + 1]; + float4 d0 = sdOV_vec4[v8 * 2], d1 = sdOV_vec4[v8 * 2 + 1]; + +#define COMP(i, sf, df) \ + float sh##i = (sf) - lse_val; \ + float p##i = (sh##i < -80.0f) ? 0.0f : __expf(sh##i); \ + float ds##i = p##i * softmax_scale * ((df) - row_dot_val); + + COMP(0, s0.x, d0.x) + COMP(1, s0.y, d0.y) + COMP(2, s0.z, d0.z) + COMP(3, s0.w, d0.w) + COMP(4, s1.x, d1.x) + COMP(5, s1.y, d1.y) + COMP(6, s1.z, d1.z) + COMP(7, s1.w, d1.w) +#undef COMP + + uint4 res; + asm volatile( + "{ mov.b32 %0, {%4,%5}; mov.b32 %1, {%6,%7}; mov.b32 %2, {%8,%9}; " + "mov.b32 %3, {%10,%11}; }\n" + : "=r"(res.x), "=r"(res.y), "=r"(res.z), "=r"(res.w) + : "h"(__half_as_ushort(__float2half_rn(ds0))), + "h"(__half_as_ushort(__float2half_rn(ds1))), + "h"(__half_as_ushort(__float2half_rn(ds2))), + "h"(__half_as_ushort(__float2half_rn(ds3))), + "h"(__half_as_ushort(__float2half_rn(ds4))), + "h"(__half_as_ushort(__float2half_rn(ds5))), + "h"(__half_as_ushort(__float2half_rn(ds6))), + "h"(__half_as_ushort(__float2half_rn(ds7)))); + buf[cnt++] = res; + } + +#pragma unroll + for (int c = tail_start + thread_in_row; c < BLOCK_N; + c += THREADS_PER_ROW) { + float s = (c < valid_k_rows) ? sS_row[c] : NEG_INF; + float dov = (c < valid_k_rows) ? sdOV_row[c] : 0.0f; + float p = (s - lse_val < -80.0f) ? 0.0f : __expf(s - lse_val); + float ds = p * softmax_scale * + ((c < valid_k_rows) ? (dov - row_dot_val) : 0.0f); + sdS_row[c] = __float2half_rn(ds); + } + +#pragma unroll + for (int i = 0; i < cnt; ++i) { + const int v8 = thread_in_row + i * THREADS_PER_ROW; + if (v8 < vec8_cols) { + sdS_vec_u4[v8] = buf[i]; } - __syncthreads(); + } } + __syncthreads(); - const int total_fp16_x4 = (valid_q_rows * D) / 4; - for (int i = tid; i < total_fp16_x4; i += THREADS_PER_BLOCK) { - const int row = i / (D / 4); - const int col = (i % (D / 4)) * 4; - - const float* s_dQ_row = sdQ + row * Q_STRIDE; - - const __half h0 = __float2half_rn(s_dQ_row[col + 0]); - const __half h1 = __float2half_rn(s_dQ_row[col + 1]); - const __half h2 = __float2half_rn(s_dQ_row[col + 2]); - const __half h3 = __float2half_rn(s_dQ_row[col + 3]); - - asm volatile( - "st.global.v4.u16 [%0], {%1, %2, %3, %4};" - : - : "l"(dQ_ptr + row * D + col), - "h"(__half_as_ushort(h0)), - "h"(__half_as_ushort(h1)), - "h"(__half_as_ushort(h2)), - "h"(__half_as_ushort(h3)) - : "memory" - ); + const int num_tiles_m_dq = (BLOCK_M + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_dq = (D + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_dq = (BLOCK_N + WMMA_K - 1) / WMMA_K; + const int total_tiles_dq = num_tiles_m_dq * num_tiles_n_dq; + const int tiles_per_warp_dq = + (total_tiles_dq + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + + for (int tile_local = 0; tile_local < tiles_per_warp_dq; ++tile_local) { + const int global_tile_idx = warp_id * tiles_per_warp_dq + tile_local; + if (global_tile_idx >= total_tiles_dq) break; + + const int tile_m_idx = global_tile_idx / num_tiles_n_dq; + const int tile_n_idx = global_tile_idx % num_tiles_n_dq; + + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; + + if (tile_m >= valid_q_rows || tile_n >= D) continue; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + fill_fragment(acc_frag, 0.0f); + +#pragma unroll + for (int k_tile = 0; k_tile < num_tiles_k_dq; ++k_tile) { + const int k_offset = k_tile * WMMA_K; + if (k_offset >= valid_k_rows) break; + + load_matrix_sync(a_frag, sdS + tile_m * S_STRIDE + k_offset, S_STRIDE); + load_matrix_sync(b_frag, sK + k_offset * KV_STRIDE + tile_n, KV_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + fragment curr_frag; + load_matrix_sync(curr_frag, sdQ + tile_m * Q_STRIDE + tile_n, Q_STRIDE, + mem_row_major); + +#pragma unroll + for (int i = 0; i < curr_frag.num_elements; ++i) { + curr_frag.x[i] += acc_frag.x[i]; + } + store_matrix_sync(sdQ + tile_m * Q_STRIDE + tile_n, curr_frag, Q_STRIDE, + mem_row_major); } + __syncthreads(); + } + + const int total_fp16_x4 = (valid_q_rows * D) / 4; + for (int i = tid; i < total_fp16_x4; i += THREADS_PER_BLOCK) { + const int row = i / (D / 4); + const int col = (i % (D / 4)) * 4; + + const float* s_dQ_row = sdQ + row * Q_STRIDE; + + const __half h0 = __float2half_rn(s_dQ_row[col + 0]); + const __half h1 = __float2half_rn(s_dQ_row[col + 1]); + const __half h2 = __float2half_rn(s_dQ_row[col + 2]); + const __half h3 = __float2half_rn(s_dQ_row[col + 3]); + + asm volatile("st.global.v4.u16 [%0], {%1, %2, %3, %4};" + : + : "l"(dQ_ptr + row * D + col), "h"(__half_as_ushort(h0)), + "h"(__half_as_ushort(h1)), "h"(__half_as_ushort(h2)), + "h"(__half_as_ushort(h3)) + : "memory"); + } } -template +template __global__ void __launch_bounds__(dKVKernelConfig::THREADS_PER_BLOCK, 2) -flash_attention_backward_dkv_kernel( - const __half* __restrict__ Q, - const __half* __restrict__ K, - const __half* __restrict__ V, - const __half* __restrict__ O, - const __half* __restrict__ dO, - const float* __restrict__ softmax_lse, - __half* __restrict__ dK, - __half* __restrict__ dV, - const int B, - const int H, - const int M, - const int N, - const float softmax_scale -) { - using Config = dKVKernelConfig; - constexpr int BLOCK_M = Config::BLOCK_M; - constexpr int BLOCK_N = Config::BLOCK_N; - constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; - constexpr int THREADS_PER_ROW = Config::THREADS_PER_ROW; - constexpr int WARPS_PER_BLOCK = Config::WARPS_PER_BLOCK; - constexpr int Q_STRIDE = Config::Q_STRIDE; - constexpr int KV_STRIDE = Config::KV_STRIDE; - constexpr int S_STRIDE = Config::S_STRIDE; - constexpr int PER_UINT4 = Config::PER_UINT4; - constexpr int NUM_UINT4_Q_BLOCK = Config::NUM_UINT4_Q_BLOCK; - constexpr int NUM_UINT4_KV_BLOCK = Config::NUM_UINT4_KV_BLOCK; - - const float NEG_INF = -1e30f; - const int batch_head_id = blockIdx.z; - if (batch_head_id >= B * H) return; - - const int block_kv = blockIdx.x; - const int start_kv = block_kv * BLOCK_M; - if (start_kv >= N) return; - - int num_q_tiles = (M + BLOCK_N - 1) / BLOCK_N; - const int valid_kv_rows = min(BLOCK_M, N - start_kv); - - const int tid = threadIdx.x; - const int warp_id = tid / MAX_THREADS_PER_WARP; - const int lane_id = tid % MAX_THREADS_PER_WARP; - - const __half* q_ptr = Q + (size_t)batch_head_id * M * D; - const __half* k_ptr = K + (size_t)batch_head_id * N * D + start_kv * D; - const __half* v_ptr = V + (size_t)batch_head_id * N * D + start_kv * D; - const __half* o_ptr = O + (size_t)batch_head_id * M * D; - const __half* dO_ptr = dO + (size_t)batch_head_id * M * D; - const float* lse_ptr = softmax_lse + (size_t)batch_head_id * M; - __half* dK_ptr = dK + (size_t)batch_head_id * N * D + start_kv * D; - __half* dV_ptr = dV + (size_t)batch_head_id * N * D + start_kv * D; - - extern __shared__ char smem_raw[]; - init_smem(smem_raw); - auto& smem = *reinterpret_cast(smem_raw); - - __half* sK = smem.k; - __half* sV = smem.v; - __half* sdO = smem.reuse_qdO.dO; - __half* sQ = smem.reuse_qdO.q; - float* sS = smem.reuse_sp.s; - __half* sP = smem.reuse_sp.p; - float* sdOV = smem.reuse_dOVS.dOV; - __half* sdS = smem.reuse_dOVS.dS; - float* sRowDot = smem.row_dot; - float* sLse = smem.lse; - float* sdK = smem.dK; - float* sdV = smem.dV; - - const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; - const int q_stride_uint4 = (Q_STRIDE + PER_UINT4 - 1) / PER_UINT4; - const int kv_stride_uint4 = (KV_STRIDE + PER_UINT4 - 1) / PER_UINT4; - - const uint4* k_vec = reinterpret_cast(k_ptr); - const uint4* v_vec = reinterpret_cast(v_ptr); - uint4* sK_vec = reinterpret_cast(sK); - uint4* sV_vec = reinterpret_cast(sV); - - #pragma unroll 2 - for (int idx = tid; idx < NUM_UINT4_KV_BLOCK; idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; + flash_attention_backward_dkv_kernel( + const __half* __restrict__ Q, const __half* __restrict__ K, + const __half* __restrict__ V, const __half* __restrict__ O, + const __half* __restrict__ dO, const float* __restrict__ softmax_lse, + __half* __restrict__ dK, __half* __restrict__ dV, const int B, + const int H, const int M, const int N, const float softmax_scale) { + using Config = dKVKernelConfig; + constexpr int BLOCK_M = Config::BLOCK_M; + constexpr int BLOCK_N = Config::BLOCK_N; + constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; + constexpr int THREADS_PER_ROW = Config::THREADS_PER_ROW; + constexpr int WARPS_PER_BLOCK = Config::WARPS_PER_BLOCK; + constexpr int Q_STRIDE = Config::Q_STRIDE; + constexpr int KV_STRIDE = Config::KV_STRIDE; + constexpr int S_STRIDE = Config::S_STRIDE; + constexpr int PER_UINT4 = Config::PER_UINT4; + constexpr int NUM_UINT4_Q_BLOCK = Config::NUM_UINT4_Q_BLOCK; + constexpr int NUM_UINT4_KV_BLOCK = Config::NUM_UINT4_KV_BLOCK; + + const float NEG_INF = -1e30f; + const int batch_head_id = blockIdx.z; + if (batch_head_id >= B * H) return; + + const int block_kv = blockIdx.x; + const int start_kv = block_kv * BLOCK_M; + if (start_kv >= N) return; + + int num_q_tiles = (M + BLOCK_N - 1) / BLOCK_N; + const int valid_kv_rows = min(BLOCK_M, N - start_kv); + + const int tid = threadIdx.x; + const int warp_id = tid / MAX_THREADS_PER_WARP; + const int lane_id = tid % MAX_THREADS_PER_WARP; + + const __half* q_ptr = Q + (size_t)batch_head_id * M * D; + const __half* k_ptr = K + (size_t)batch_head_id * N * D + start_kv * D; + const __half* v_ptr = V + (size_t)batch_head_id * N * D + start_kv * D; + const __half* o_ptr = O + (size_t)batch_head_id * M * D; + const __half* dO_ptr = dO + (size_t)batch_head_id * M * D; + const float* lse_ptr = softmax_lse + (size_t)batch_head_id * M; + __half* dK_ptr = dK + (size_t)batch_head_id * N * D + start_kv * D; + __half* dV_ptr = dV + (size_t)batch_head_id * N * D + start_kv * D; + + extern __shared__ char smem_raw[]; + init_smem(smem_raw); + auto& smem = *reinterpret_cast(smem_raw); + + __half* sK = smem.k; + __half* sV = smem.v; + __half* sdO = smem.reuse_qdO.dO; + __half* sQ = smem.reuse_qdO.q; + float* sS = smem.reuse_sp.s; + __half* sP = smem.reuse_sp.p; + float* sdOV = smem.reuse_dOVS.dOV; + __half* sdS = smem.reuse_dOVS.dS; + float* sRowDot = smem.row_dot; + float* sLse = smem.lse; + float* sdK = smem.dK; + float* sdV = smem.dV; + + const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; + const int q_stride_uint4 = (Q_STRIDE + PER_UINT4 - 1) / PER_UINT4; + const int kv_stride_uint4 = (KV_STRIDE + PER_UINT4 - 1) / PER_UINT4; + + const uint4* k_vec = reinterpret_cast(k_ptr); + const uint4* v_vec = reinterpret_cast(v_ptr); + uint4* sK_vec = reinterpret_cast(sK); + uint4* sV_vec = reinterpret_cast(sV); + +#pragma unroll 2 + for (int idx = tid; idx < NUM_UINT4_KV_BLOCK; idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + + uint4 k_val = make_uint4(0, 0, 0, 0); + uint4 v_val = make_uint4(0, 0, 0, 0); + + if (row < valid_kv_rows && vec_col < d_stride_uint4) { + k_val = __ldg(&k_vec[row * d_stride_uint4 + vec_col]); + v_val = __ldg(&v_vec[row * d_stride_uint4 + vec_col]); + } + sK_vec[row * kv_stride_uint4 + vec_col] = k_val; + sV_vec[row * kv_stride_uint4 + vec_col] = v_val; + } + __syncthreads(); - uint4 k_val = make_uint4(0, 0, 0, 0); - uint4 v_val = make_uint4(0, 0, 0, 0); + for (int block_n = 0; block_n < num_q_tiles; ++block_n) { + const int start_col = block_n * BLOCK_N; + if (start_col >= M) break; + const int valid_q_rows = min(BLOCK_N, M - start_col); - if (row < valid_kv_rows && vec_col < d_stride_uint4) { - k_val = __ldg(&k_vec[row * d_stride_uint4 + vec_col]); - v_val = __ldg(&v_vec[row * d_stride_uint4 + vec_col]); - } - sK_vec[row * kv_stride_uint4 + vec_col] = k_val; - sV_vec[row * kv_stride_uint4 + vec_col] = v_val; + if constexpr (IS_CAUSAL) { + if (start_kv >= start_col + valid_q_rows) { + continue; + } } - __syncthreads(); - for (int block_n = 0; block_n < num_q_tiles; ++block_n) { - const int start_col = block_n * BLOCK_N; - if (start_col >= M) break; - const int valid_q_rows = min(BLOCK_N, M - start_col); - - if constexpr (IS_CAUSAL) { - if (start_kv >= start_col + valid_q_rows) { continue; } - } + const uint4* q_vec = reinterpret_cast(q_ptr + start_col * D); + uint4* sQ_vec = reinterpret_cast(sdO); - const uint4* q_vec = reinterpret_cast(q_ptr + start_col * D); - uint4* sQ_vec = reinterpret_cast(sdO); - - #pragma unroll 2 - for (int idx = tid; idx < NUM_UINT4_Q_BLOCK; idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - uint4 q_val = make_uint4(0, 0, 0, 0); - if (row < valid_q_rows && vec_col < d_stride_uint4) { - q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); - } - sQ_vec[row * q_stride_uint4 + vec_col] = q_val; - } - __syncthreads(); - - const int num_tiles_m_qk = (BLOCK_N + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_qk = (BLOCK_M + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_qk = (D + WMMA_K - 1) / WMMA_K; - const int total_tiles_qk = num_tiles_m_qk * num_tiles_n_qk; - const int tiles_per_warp_qk = (total_tiles_qk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - const unsigned row_causal = (lane_id & 0b1) + ((lane_id >> 2) & 0b1) * 8 + ((lane_id >> 4) & 0b1) * 4; - const unsigned col_causal = ((lane_id >> 1) & 0b1) * 2 + ((lane_id >> 3) & 0b1) * 8; - - for (int tile_local = 0; tile_local < tiles_per_warp_qk; ++tile_local) { - const int tile_idx = warp_id * tiles_per_warp_qk + tile_local; - if (tile_idx >= total_tiles_qk) break; - - const int tile_m_idx = tile_idx / num_tiles_n_qk; - const int tile_n_idx = tile_idx % num_tiles_n_qk; - - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; - - if (tile_m >= valid_q_rows || tile_n >= valid_kv_rows) continue; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - fill_fragment(acc_frag, 0.0f); - - #pragma unroll - for (int k_tile = 0; k_tile < num_tiles_k_qk; ++k_tile) { - const int k_offset = k_tile * WMMA_K; - if (k_offset >= D) break; - load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, Q_STRIDE); - load_matrix_sync(b_frag, sK + tile_n * KV_STRIDE + k_offset, KV_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - - if constexpr (IS_CAUSAL) { - #pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - const unsigned col = col_causal + (i & 0b1) + ((i >> 2) & 0b1) * 4; - const unsigned row = row_causal + ((i >> 1) & 0b1) * 2; - - const int global_m = start_col + tile_m + row; - const int global_n = start_kv + tile_n + col; - - const bool is_valid = (global_m < start_col + valid_q_rows) && - (global_n < start_kv + valid_kv_rows); - - acc_frag.x[i] = is_valid - ? ((global_n > global_m) ? NEG_INF : acc_frag.x[i] * softmax_scale) - : NEG_INF; - } - } else { - #pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - acc_frag.x[i] *= softmax_scale; - } - } - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, mem_row_major); - } - __syncthreads(); - - const uint4* do_vec = reinterpret_cast(dO_ptr + start_col * D); - uint4* sdO_vec = reinterpret_cast(sdO); - - #pragma unroll 2 - for (int idx = tid; idx < NUM_UINT4_Q_BLOCK; idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - uint4 do_val = make_uint4(0, 0, 0, 0); - if (row < valid_q_rows && vec_col < d_stride_uint4) { - do_val = __ldg(&do_vec[row * d_stride_uint4 + vec_col]); - } - sdO_vec[row * q_stride_uint4 + vec_col] = do_val; - } - __syncthreads(); - - const __half* current_o_ptr = o_ptr + start_col * D; - - if (tid < valid_q_rows * THREADS_PER_ROW) { - const int row = tid / THREADS_PER_ROW; - const int thread_in_row = tid % THREADS_PER_ROW; - const int fp16_x4_per_row = D / 4; - const int work_per_thread = (fp16_x4_per_row + THREADS_PER_ROW - 1) / THREADS_PER_ROW; - const unsigned mask = (valid_q_rows == BLOCK_N) ? 0xFFFFFFFFU : __activemask(); - - float thread_dot = 0.0f; - - #pragma unroll - for (int j = 0; j < work_per_thread; ++j) { - const int chunk_idx = thread_in_row + j * THREADS_PER_ROW; - if (chunk_idx >= fp16_x4_per_row) break; - const int col = chunk_idx * 4; - - const half* o_addr = current_o_ptr + row * D + col; - ushort o_h0, o_h1, o_h2, o_h3; - asm volatile( - "ld.global.v4.u16 {%0, %1, %2, %3}, [%4];" - : "=h"(o_h0), "=h"(o_h1), "=h"(o_h2), "=h"(o_h3) - : "l"(o_addr) - : "memory" - ); - - const half* dO_addr = sdO + row * Q_STRIDE + col; - ushort d_h0, d_h1, d_h2, d_h3; - const uint32_t ptr_dO = static_cast(reinterpret_cast(__cvta_generic_to_shared(dO_addr))); - asm volatile( - "ld.shared.v4.u16 {%0, %1, %2, %3}, [%4];" - : "=h"(d_h0), "=h"(d_h1), "=h"(d_h2), "=h"(d_h3) - : "r"(ptr_dO) - : "memory" - ); - - const float fo_0 = __half2float(__ushort_as_half(o_h0)); - const float fo_1 = __half2float(__ushort_as_half(o_h1)); - const float fo_2 = __half2float(__ushort_as_half(o_h2)); - const float fo_3 = __half2float(__ushort_as_half(o_h3)); - - const float fd_0 = __half2float(__ushort_as_half(d_h0)); - const float fd_1 = __half2float(__ushort_as_half(d_h1)); - const float fd_2 = __half2float(__ushort_as_half(d_h2)); - const float fd_3 = __half2float(__ushort_as_half(d_h3)); - - thread_dot = __fmaf_rn(fo_0, fd_0, thread_dot); - thread_dot = __fmaf_rn(fo_1, fd_1, thread_dot); - thread_dot = __fmaf_rn(fo_2, fd_2, thread_dot); - thread_dot = __fmaf_rn(fo_3, fd_3, thread_dot); - } - - #pragma unroll - for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) - thread_dot += __shfl_down_sync(mask, thread_dot, o, THREADS_PER_ROW); - - if (thread_in_row == 0) { sRowDot[row] = thread_dot; } - } +#pragma unroll 2 + for (int idx = tid; idx < NUM_UINT4_Q_BLOCK; idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + uint4 q_val = make_uint4(0, 0, 0, 0); + if (row < valid_q_rows && vec_col < d_stride_uint4) { + q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); + } + sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + } + __syncthreads(); - if (tid < valid_q_rows) { sLse[tid] = lse_ptr[start_col + tid]; } - __syncthreads(); - - const int num_tiles_m_dov = (BLOCK_N + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_dov = (BLOCK_M + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_dov = (D + WMMA_K - 1) / WMMA_K; - const int total_tiles_dov = num_tiles_m_dov * num_tiles_n_dov; - const int tiles_per_warp_dov = (total_tiles_dov + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - - for (int tile_local = 0; tile_local < tiles_per_warp_dov; ++tile_local) { - const int tile_idx = warp_id * tiles_per_warp_dov + tile_local; - if (tile_idx >= total_tiles_dov) break; - - const int tile_m_idx = tile_idx / num_tiles_n_dov; - const int tile_n_idx = tile_idx % num_tiles_n_dov; - - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; - - if (tile_m >= valid_q_rows || tile_n >= valid_kv_rows) continue; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - fill_fragment(acc_frag, 0.0f); - - #pragma unroll - for (int k_tile = 0; k_tile < num_tiles_k_dov; ++k_tile) { - const int k_offset = k_tile * WMMA_K; - if (k_offset >= D) break; - load_matrix_sync(a_frag, sdO + tile_m * Q_STRIDE + k_offset, Q_STRIDE); - load_matrix_sync(b_frag, sV + tile_n * KV_STRIDE + k_offset, KV_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - store_matrix_sync(sdOV + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, mem_row_major); + const int num_tiles_m_qk = (BLOCK_N + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_qk = (BLOCK_M + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_qk = (D + WMMA_K - 1) / WMMA_K; + const int total_tiles_qk = num_tiles_m_qk * num_tiles_n_qk; + const int tiles_per_warp_qk = + (total_tiles_qk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + const unsigned row_causal = (lane_id & 0b1) + ((lane_id >> 2) & 0b1) * 8 + + ((lane_id >> 4) & 0b1) * 4; + const unsigned col_causal = + ((lane_id >> 1) & 0b1) * 2 + ((lane_id >> 3) & 0b1) * 8; + + for (int tile_local = 0; tile_local < tiles_per_warp_qk; ++tile_local) { + const int tile_idx = warp_id * tiles_per_warp_qk + tile_local; + if (tile_idx >= total_tiles_qk) break; + + const int tile_m_idx = tile_idx / num_tiles_n_qk; + const int tile_n_idx = tile_idx % num_tiles_n_qk; + + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; + + if (tile_m >= valid_q_rows || tile_n >= valid_kv_rows) continue; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + fill_fragment(acc_frag, 0.0f); + +#pragma unroll + for (int k_tile = 0; k_tile < num_tiles_k_qk; ++k_tile) { + const int k_offset = k_tile * WMMA_K; + if (k_offset >= D) break; + load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, Q_STRIDE); + load_matrix_sync(b_frag, sK + tile_n * KV_STRIDE + k_offset, KV_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + if constexpr (IS_CAUSAL) { +#pragma unroll + for (int i = 0; i < acc_frag.num_elements; ++i) { + const unsigned col = col_causal + (i & 0b1) + ((i >> 2) & 0b1) * 4; + const unsigned row = row_causal + ((i >> 1) & 0b1) * 2; + + const int global_m = start_col + tile_m + row; + const int global_n = start_kv + tile_n + col; + + const bool is_valid = (global_m < start_col + valid_q_rows) && + (global_n < start_kv + valid_kv_rows); + + acc_frag.x[i] = + is_valid ? ((global_n > global_m) ? NEG_INF + : acc_frag.x[i] * softmax_scale) + : NEG_INF; } - __syncthreads(); - - const int total_elements = BLOCK_N * BLOCK_M; - const int total_pairs = (total_elements + 1) / 2; - - #pragma unroll 2 - for (int i = tid; i < total_pairs; i += THREADS_PER_BLOCK) { - const int linear_idx0 = i * 2; - const int linear_idx1 = linear_idx0 + 1; - - const int row0 = linear_idx0 / BLOCK_M; - const int col0 = linear_idx0 % BLOCK_M; - const bool has_pair = (linear_idx1 < total_elements); - const int row1 = has_pair ? linear_idx1 / BLOCK_M : row0; - const int col1 = has_pair ? linear_idx1 % BLOCK_M : col0 + 1; - - float s0 = 0.0f, s1 = 0.0f; - float dov0 = 0.0f, dov1 = 0.0f; - - const bool in_bounds0 = (row0 < valid_q_rows) && (col0 < valid_kv_rows); - const bool causal_ok0 = !IS_CAUSAL || ((start_kv + col0) <= (start_col + row0)); - const bool valid0 = in_bounds0 && causal_ok0; - - bool valid1 = false; - if (has_pair) { - const bool in_bounds1 = (row1 < valid_q_rows) && (col1 < valid_kv_rows); - const bool causal_ok1 = !IS_CAUSAL || ((start_kv + col1) <= (start_col + row1)); - valid1 = in_bounds1 && causal_ok1; - } - - if (valid0) { s0 = sS[row0 * S_STRIDE + col0]; dov0 = sdOV[row0 * S_STRIDE + col0]; } else { s0 = NEG_INF; } - if (valid1) { s1 = sS[row1 * S_STRIDE + col1]; dov1 = sdOV[row1 * S_STRIDE + col1]; } else { s1 = NEG_INF; } - - float lse0 = sLse[row0]; - float lse1 = (valid1 && row1 != row0) ? sLse[row1] : lse0; - float row_dot0 = sRowDot[row0]; - float row_dot1 = (valid1 && row1 != row0) ? sRowDot[row1] : row_dot0; - - float shifted0 = s0 - lse0; - float shifted1 = s1 - lse1; - - float p0 = (shifted0 < -80.0f) ? 0.0f : __expf(shifted0); - float p1 = (shifted1 < -80.0f) ? 0.0f : __expf(shifted1); - - float diff0 = dov0 - row_dot0; - float diff1 = dov1 - row_dot1; - float ds0 = valid0 ? fmaf(p0, softmax_scale * diff0, 0.0f) : 0.0f; - float ds1 = valid1 ? fmaf(p1, softmax_scale * diff1, 0.0f) : 0.0f; - - __half2 p_h2 = __float22half2_rn(make_float2(p0, p1)); - __half2 ds_h2 = __float22half2_rn(make_float2(ds0, ds1)); - - if (col1 < BLOCK_M && row1 == row0) { - __half* p_dst = sP + row0 * BLOCK_M + col0; - __half* ds_dst = sdS + row0 * BLOCK_M + col0; - p_dst[0] = p_h2.x; p_dst[1] = p_h2.y; - ds_dst[0] = ds_h2.x; ds_dst[1] = ds_h2.y; - } else { - if (valid0) { sP[row0 * BLOCK_M + col0] = p_h2.x; sdS[row0 * BLOCK_M + col0] = ds_h2.x; } - if (valid1) { sP[row1 * BLOCK_M + col1] = p_h2.y; sdS[row1 * BLOCK_M + col1] = ds_h2.y; } - } + } else { +#pragma unroll + for (int i = 0; i < acc_frag.num_elements; ++i) { + acc_frag.x[i] *= softmax_scale; } - __syncthreads(); - - const int num_tiles_m_dv = (BLOCK_M + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_dv = (D + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_dv = (BLOCK_N + WMMA_K - 1) / WMMA_K; - const int total_tiles_dv = num_tiles_m_dv * num_tiles_n_dv; - const int tiles_per_warp_dv = (total_tiles_dv + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + } + store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, + mem_row_major); + } + __syncthreads(); - for (int tile_local = 0; tile_local < tiles_per_warp_dv; ++tile_local) { - const int tile_idx = warp_id * tiles_per_warp_dv + tile_local; - if (tile_idx >= total_tiles_dv) break; + const uint4* do_vec = + reinterpret_cast(dO_ptr + start_col * D); + uint4* sdO_vec = reinterpret_cast(sdO); - const int tile_m_idx = tile_idx / num_tiles_n_dv; - const int tile_n_idx = tile_idx % num_tiles_n_dv; +#pragma unroll 2 + for (int idx = tid; idx < NUM_UINT4_Q_BLOCK; idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + uint4 do_val = make_uint4(0, 0, 0, 0); + if (row < valid_q_rows && vec_col < d_stride_uint4) { + do_val = __ldg(&do_vec[row * d_stride_uint4 + vec_col]); + } + sdO_vec[row * q_stride_uint4 + vec_col] = do_val; + } + __syncthreads(); - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; + const __half* current_o_ptr = o_ptr + start_col * D; - if (tile_m >= valid_kv_rows || tile_n >= D) continue; + if (tid < valid_q_rows * THREADS_PER_ROW) { + const int row = tid / THREADS_PER_ROW; + const int thread_in_row = tid % THREADS_PER_ROW; + const int fp16_x4_per_row = D / 4; + const int work_per_thread = + (fp16_x4_per_row + THREADS_PER_ROW - 1) / THREADS_PER_ROW; + const unsigned mask = + (valid_q_rows == BLOCK_N) ? 0xFFFFFFFFU : __activemask(); + + float thread_dot = 0.0f; + +#pragma unroll + for (int j = 0; j < work_per_thread; ++j) { + const int chunk_idx = thread_in_row + j * THREADS_PER_ROW; + if (chunk_idx >= fp16_x4_per_row) break; + const int col = chunk_idx * 4; + + const half* o_addr = current_o_ptr + row * D + col; + ushort o_h0, o_h1, o_h2, o_h3; + asm volatile("ld.global.v4.u16 {%0, %1, %2, %3}, [%4];" + : "=h"(o_h0), "=h"(o_h1), "=h"(o_h2), "=h"(o_h3) + : "l"(o_addr) + : "memory"); + + const half* dO_addr = sdO + row * Q_STRIDE + col; + ushort d_h0, d_h1, d_h2, d_h3; + const uint32_t ptr_dO = static_cast( + reinterpret_cast(__cvta_generic_to_shared(dO_addr))); + asm volatile("ld.shared.v4.u16 {%0, %1, %2, %3}, [%4];" + : "=h"(d_h0), "=h"(d_h1), "=h"(d_h2), "=h"(d_h3) + : "r"(ptr_dO) + : "memory"); + + const float fo_0 = __half2float(__ushort_as_half(o_h0)); + const float fo_1 = __half2float(__ushort_as_half(o_h1)); + const float fo_2 = __half2float(__ushort_as_half(o_h2)); + const float fo_3 = __half2float(__ushort_as_half(o_h3)); + + const float fd_0 = __half2float(__ushort_as_half(d_h0)); + const float fd_1 = __half2float(__ushort_as_half(d_h1)); + const float fd_2 = __half2float(__ushort_as_half(d_h2)); + const float fd_3 = __half2float(__ushort_as_half(d_h3)); + + thread_dot = __fmaf_rn(fo_0, fd_0, thread_dot); + thread_dot = __fmaf_rn(fo_1, fd_1, thread_dot); + thread_dot = __fmaf_rn(fo_2, fd_2, thread_dot); + thread_dot = __fmaf_rn(fo_3, fd_3, thread_dot); + } + +#pragma unroll + for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) + thread_dot += __shfl_down_sync(mask, thread_dot, o, THREADS_PER_ROW); + + if (thread_in_row == 0) { + sRowDot[row] = thread_dot; + } + } - fragment a_frag; - fragment b_frag; - fragment acc_frag; - load_matrix_sync(acc_frag, sdV + tile_m * KV_STRIDE + tile_n, KV_STRIDE, mem_row_major); + if (tid < valid_q_rows) { + sLse[tid] = lse_ptr[start_col + tid]; + } + __syncthreads(); - #pragma unroll - for (int k_tile = 0; k_tile < num_tiles_k_dv; ++k_tile) { - const int k_offset = k_tile * WMMA_K; - if (k_offset >= valid_q_rows) break; + const int num_tiles_m_dov = (BLOCK_N + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_dov = (BLOCK_M + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_dov = (D + WMMA_K - 1) / WMMA_K; + const int total_tiles_dov = num_tiles_m_dov * num_tiles_n_dov; + const int tiles_per_warp_dov = + (total_tiles_dov + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + + for (int tile_local = 0; tile_local < tiles_per_warp_dov; ++tile_local) { + const int tile_idx = warp_id * tiles_per_warp_dov + tile_local; + if (tile_idx >= total_tiles_dov) break; + + const int tile_m_idx = tile_idx / num_tiles_n_dov; + const int tile_n_idx = tile_idx % num_tiles_n_dov; + + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; + + if (tile_m >= valid_q_rows || tile_n >= valid_kv_rows) continue; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + fill_fragment(acc_frag, 0.0f); + +#pragma unroll + for (int k_tile = 0; k_tile < num_tiles_k_dov; ++k_tile) { + const int k_offset = k_tile * WMMA_K; + if (k_offset >= D) break; + load_matrix_sync(a_frag, sdO + tile_m * Q_STRIDE + k_offset, Q_STRIDE); + load_matrix_sync(b_frag, sV + tile_n * KV_STRIDE + k_offset, KV_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + store_matrix_sync(sdOV + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, + mem_row_major); + } + __syncthreads(); - load_matrix_sync(a_frag, sP + k_offset * BLOCK_M + tile_m, BLOCK_M); - load_matrix_sync(b_frag, sdO + k_offset * Q_STRIDE + tile_n, Q_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - store_matrix_sync(sdV + tile_m * KV_STRIDE + tile_n, acc_frag, KV_STRIDE, mem_row_major); + const int total_elements = BLOCK_N * BLOCK_M; + const int total_pairs = (total_elements + 1) / 2; + +#pragma unroll 2 + for (int i = tid; i < total_pairs; i += THREADS_PER_BLOCK) { + const int linear_idx0 = i * 2; + const int linear_idx1 = linear_idx0 + 1; + + const int row0 = linear_idx0 / BLOCK_M; + const int col0 = linear_idx0 % BLOCK_M; + const bool has_pair = (linear_idx1 < total_elements); + const int row1 = has_pair ? linear_idx1 / BLOCK_M : row0; + const int col1 = has_pair ? linear_idx1 % BLOCK_M : col0 + 1; + + float s0 = 0.0f, s1 = 0.0f; + float dov0 = 0.0f, dov1 = 0.0f; + + const bool in_bounds0 = (row0 < valid_q_rows) && (col0 < valid_kv_rows); + const bool causal_ok0 = + !IS_CAUSAL || ((start_kv + col0) <= (start_col + row0)); + const bool valid0 = in_bounds0 && causal_ok0; + + bool valid1 = false; + if (has_pair) { + const bool in_bounds1 = (row1 < valid_q_rows) && (col1 < valid_kv_rows); + const bool causal_ok1 = + !IS_CAUSAL || ((start_kv + col1) <= (start_col + row1)); + valid1 = in_bounds1 && causal_ok1; + } + + if (valid0) { + s0 = sS[row0 * S_STRIDE + col0]; + dov0 = sdOV[row0 * S_STRIDE + col0]; + } else { + s0 = NEG_INF; + } + if (valid1) { + s1 = sS[row1 * S_STRIDE + col1]; + dov1 = sdOV[row1 * S_STRIDE + col1]; + } else { + s1 = NEG_INF; + } + + float lse0 = sLse[row0]; + float lse1 = (valid1 && row1 != row0) ? sLse[row1] : lse0; + float row_dot0 = sRowDot[row0]; + float row_dot1 = (valid1 && row1 != row0) ? sRowDot[row1] : row_dot0; + + float shifted0 = s0 - lse0; + float shifted1 = s1 - lse1; + + float p0 = (shifted0 < -80.0f) ? 0.0f : __expf(shifted0); + float p1 = (shifted1 < -80.0f) ? 0.0f : __expf(shifted1); + + float diff0 = dov0 - row_dot0; + float diff1 = dov1 - row_dot1; + float ds0 = valid0 ? fmaf(p0, softmax_scale * diff0, 0.0f) : 0.0f; + float ds1 = valid1 ? fmaf(p1, softmax_scale * diff1, 0.0f) : 0.0f; + + __half2 p_h2 = __float22half2_rn(make_float2(p0, p1)); + __half2 ds_h2 = __float22half2_rn(make_float2(ds0, ds1)); + + if (col1 < BLOCK_M && row1 == row0) { + __half* p_dst = sP + row0 * BLOCK_M + col0; + __half* ds_dst = sdS + row0 * BLOCK_M + col0; + p_dst[0] = p_h2.x; + p_dst[1] = p_h2.y; + ds_dst[0] = ds_h2.x; + ds_dst[1] = ds_h2.y; + } else { + if (valid0) { + sP[row0 * BLOCK_M + col0] = p_h2.x; + sdS[row0 * BLOCK_M + col0] = ds_h2.x; } - __syncthreads(); - - __half* sQ = sdO; - #pragma unroll 2 - for (int idx = tid; idx < NUM_UINT4_Q_BLOCK; idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - uint4 q_val = make_uint4(0, 0, 0, 0); - if (row < valid_q_rows && vec_col < d_stride_uint4) { - q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); - } - sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + if (valid1) { + sP[row1 * BLOCK_M + col1] = p_h2.y; + sdS[row1 * BLOCK_M + col1] = ds_h2.y; } - __syncthreads(); - - const int num_tiles_m_dk = (BLOCK_M + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_dk = (D + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_dk = (BLOCK_N + WMMA_K - 1) / WMMA_K; - const int total_tiles_dk = num_tiles_m_dk * num_tiles_n_dk; - const int tiles_per_warp_dk = (total_tiles_dk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - - for (int tile_local = 0; tile_local < tiles_per_warp_dk; ++tile_local) { - const int tile_idx = warp_id * tiles_per_warp_dk + tile_local; - if (tile_idx >= total_tiles_dk) break; - - const int tile_m_idx = tile_idx / num_tiles_n_dk; - const int tile_n_idx = tile_idx % num_tiles_n_dk; - - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; - - if (tile_m >= valid_kv_rows || tile_n >= D) continue; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - load_matrix_sync(acc_frag, sdK + tile_m * KV_STRIDE + tile_n, KV_STRIDE, mem_row_major); - - #pragma unroll - for (int k_tile = 0; k_tile < num_tiles_k_dk; ++k_tile) { - const int k_offset = k_tile * WMMA_K; - if (k_offset >= valid_q_rows) break; - - load_matrix_sync(a_frag, sdS + k_offset * BLOCK_M + tile_m, BLOCK_M); - load_matrix_sync(b_frag, sQ + k_offset * Q_STRIDE + tile_n, Q_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - store_matrix_sync(sdK + tile_m * KV_STRIDE + tile_n, acc_frag, KV_STRIDE, mem_row_major); - } - __syncthreads(); + } } + __syncthreads(); - const int total_fp16_x4 = (valid_kv_rows * D) / 4; - for (int i = tid; i < total_fp16_x4; i += THREADS_PER_BLOCK) { - const int row = i / (D / 4); - const int col = (i % (D / 4)) * 4; - - const float* s_dK_row = sdK + row * KV_STRIDE; - const float* s_dV_row = sdV + row * KV_STRIDE; - - const __half hk0 = __float2half_rn(s_dK_row[col + 0]); - const __half hk1 = __float2half_rn(s_dK_row[col + 1]); - const __half hk2 = __float2half_rn(s_dK_row[col + 2]); - const __half hk3 = __float2half_rn(s_dK_row[col + 3]); - - const __half hv0 = __float2half_rn(s_dV_row[col + 0]); - const __half hv1 = __float2half_rn(s_dV_row[col + 1]); - const __half hv2 = __float2half_rn(s_dV_row[col + 2]); - const __half hv3 = __float2half_rn(s_dV_row[col + 3]); + const int num_tiles_m_dv = (BLOCK_M + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_dv = (D + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_dv = (BLOCK_N + WMMA_K - 1) / WMMA_K; + const int total_tiles_dv = num_tiles_m_dv * num_tiles_n_dv; + const int tiles_per_warp_dv = + (total_tiles_dv + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + + for (int tile_local = 0; tile_local < tiles_per_warp_dv; ++tile_local) { + const int tile_idx = warp_id * tiles_per_warp_dv + tile_local; + if (tile_idx >= total_tiles_dv) break; + + const int tile_m_idx = tile_idx / num_tiles_n_dv; + const int tile_n_idx = tile_idx % num_tiles_n_dv; + + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; + + if (tile_m >= valid_kv_rows || tile_n >= D) continue; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + load_matrix_sync(acc_frag, sdV + tile_m * KV_STRIDE + tile_n, KV_STRIDE, + mem_row_major); + +#pragma unroll + for (int k_tile = 0; k_tile < num_tiles_k_dv; ++k_tile) { + const int k_offset = k_tile * WMMA_K; + if (k_offset >= valid_q_rows) break; + + load_matrix_sync(a_frag, sP + k_offset * BLOCK_M + tile_m, BLOCK_M); + load_matrix_sync(b_frag, sdO + k_offset * Q_STRIDE + tile_n, Q_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + store_matrix_sync(sdV + tile_m * KV_STRIDE + tile_n, acc_frag, KV_STRIDE, + mem_row_major); + } + __syncthreads(); - asm volatile( - "st.global.v4.u16 [%0], {%1, %2, %3, %4};" - : - : "l"(dK_ptr + row * D + col), - "h"(__half_as_ushort(hk0)), - "h"(__half_as_ushort(hk1)), - "h"(__half_as_ushort(hk2)), - "h"(__half_as_ushort(hk3)) - : "memory" - ); + __half* sQ = sdO; +#pragma unroll 2 + for (int idx = tid; idx < NUM_UINT4_Q_BLOCK; idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + uint4 q_val = make_uint4(0, 0, 0, 0); + if (row < valid_q_rows && vec_col < d_stride_uint4) { + q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); + } + sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + } + __syncthreads(); - asm volatile( - "st.global.v4.u16 [%0], {%1, %2, %3, %4};" - : - : "l"(dV_ptr + row * D + col), - "h"(__half_as_ushort(hv0)), - "h"(__half_as_ushort(hv1)), - "h"(__half_as_ushort(hv2)), - "h"(__half_as_ushort(hv3)) - : "memory" - ); + const int num_tiles_m_dk = (BLOCK_M + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_dk = (D + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_dk = (BLOCK_N + WMMA_K - 1) / WMMA_K; + const int total_tiles_dk = num_tiles_m_dk * num_tiles_n_dk; + const int tiles_per_warp_dk = + (total_tiles_dk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + + for (int tile_local = 0; tile_local < tiles_per_warp_dk; ++tile_local) { + const int tile_idx = warp_id * tiles_per_warp_dk + tile_local; + if (tile_idx >= total_tiles_dk) break; + + const int tile_m_idx = tile_idx / num_tiles_n_dk; + const int tile_n_idx = tile_idx % num_tiles_n_dk; + + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; + + if (tile_m >= valid_kv_rows || tile_n >= D) continue; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + load_matrix_sync(acc_frag, sdK + tile_m * KV_STRIDE + tile_n, KV_STRIDE, + mem_row_major); + +#pragma unroll + for (int k_tile = 0; k_tile < num_tiles_k_dk; ++k_tile) { + const int k_offset = k_tile * WMMA_K; + if (k_offset >= valid_q_rows) break; + + load_matrix_sync(a_frag, sdS + k_offset * BLOCK_M + tile_m, BLOCK_M); + load_matrix_sync(b_frag, sQ + k_offset * Q_STRIDE + tile_n, Q_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + store_matrix_sync(sdK + tile_m * KV_STRIDE + tile_n, acc_frag, KV_STRIDE, + mem_row_major); } + __syncthreads(); + } + + const int total_fp16_x4 = (valid_kv_rows * D) / 4; + for (int i = tid; i < total_fp16_x4; i += THREADS_PER_BLOCK) { + const int row = i / (D / 4); + const int col = (i % (D / 4)) * 4; + + const float* s_dK_row = sdK + row * KV_STRIDE; + const float* s_dV_row = sdV + row * KV_STRIDE; + + const __half hk0 = __float2half_rn(s_dK_row[col + 0]); + const __half hk1 = __float2half_rn(s_dK_row[col + 1]); + const __half hk2 = __float2half_rn(s_dK_row[col + 2]); + const __half hk3 = __float2half_rn(s_dK_row[col + 3]); + + const __half hv0 = __float2half_rn(s_dV_row[col + 0]); + const __half hv1 = __float2half_rn(s_dV_row[col + 1]); + const __half hv2 = __float2half_rn(s_dV_row[col + 2]); + const __half hv3 = __float2half_rn(s_dV_row[col + 3]); + + asm volatile("st.global.v4.u16 [%0], {%1, %2, %3, %4};" + : + : "l"(dK_ptr + row * D + col), "h"(__half_as_ushort(hk0)), + "h"(__half_as_ushort(hk1)), "h"(__half_as_ushort(hk2)), + "h"(__half_as_ushort(hk3)) + : "memory"); + + asm volatile("st.global.v4.u16 [%0], {%1, %2, %3, %4};" + : + : "l"(dV_ptr + row * D + col), "h"(__half_as_ushort(hv0)), + "h"(__half_as_ushort(hv1)), "h"(__half_as_ushort(hv2)), + "h"(__half_as_ushort(hv3)) + : "memory"); + } } -template +template void launcher_flash_attention_backward_dq( - const torch::Tensor& Q, - const torch::Tensor& K, - const torch::Tensor& V, - const torch::Tensor& O, - torch::Tensor& dO, - const torch::Tensor& softmax_lse, - torch::Tensor& dQ, - float softmax_scale, - bool is_causal, - cudaStream_t stream -) { - using Config = dQKernelConfig; - - const int B = Q.size(0); - const int H = Q.size(1); - const int M = Q.size(2); - const int N = K.size(2); - - const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; - const dim3 grid(grid_x, 1, B * H); - const dim3 block(Config::THREADS_PER_BLOCK); - const size_t smem = Config::TOTAL_SMEM; - - TORCH_CHECK(smem <= MAX_SMEM_PER_SM, "Shared memory exceeds 96KB for dQ: ", smem, " bytes"); - - auto kernel = is_causal ? - (void*)flash_attention_backward_dq_kernel : - (void*)flash_attention_backward_dq_kernel; - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - - if (is_causal) { - flash_attention_backward_dq_kernel<<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K.data_ptr()), - reinterpret_cast(V.data_ptr()), - reinterpret_cast(O.data_ptr()), - reinterpret_cast<__half*>(dO.data_ptr()), - softmax_lse.data_ptr(), - reinterpret_cast<__half*>(dQ.data_ptr()), - B, H, M, N, softmax_scale - ); - } else { - flash_attention_backward_dq_kernel<<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K.data_ptr()), - reinterpret_cast(V.data_ptr()), - reinterpret_cast(O.data_ptr()), - reinterpret_cast<__half*>(dO.data_ptr()), - softmax_lse.data_ptr(), - reinterpret_cast<__half*>(dQ.data_ptr()), - B, H, M, N, softmax_scale - ); - } + const torch::Tensor& Q, const torch::Tensor& K, const torch::Tensor& V, + const torch::Tensor& O, torch::Tensor& dO, const torch::Tensor& softmax_lse, + torch::Tensor& dQ, float softmax_scale, bool is_causal, + cudaStream_t stream) { + using Config = dQKernelConfig; + + const int B = Q.size(0); + const int H = Q.size(1); + const int M = Q.size(2); + const int N = K.size(2); + + const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; + const dim3 grid(grid_x, 1, B * H); + const dim3 block(Config::THREADS_PER_BLOCK); + const size_t smem = Config::TOTAL_SMEM; + + TORCH_CHECK(smem <= MAX_SMEM_PER_SM, + "Shared memory exceeds 96KB for dQ: ", smem, " bytes"); + + auto kernel = is_causal ? (void*)flash_attention_backward_dq_kernel + : (void*)flash_attention_backward_dq_kernel; + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + if (is_causal) { + flash_attention_backward_dq_kernel<<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K.data_ptr()), + reinterpret_cast(V.data_ptr()), + reinterpret_cast(O.data_ptr()), + reinterpret_cast<__half*>(dO.data_ptr()), softmax_lse.data_ptr(), + reinterpret_cast<__half*>(dQ.data_ptr()), B, H, M, N, softmax_scale); + } else { + flash_attention_backward_dq_kernel<<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K.data_ptr()), + reinterpret_cast(V.data_ptr()), + reinterpret_cast(O.data_ptr()), + reinterpret_cast<__half*>(dO.data_ptr()), softmax_lse.data_ptr(), + reinterpret_cast<__half*>(dQ.data_ptr()), B, H, M, N, softmax_scale); + } } -template +template void launcher_flash_attention_backward_dkv( - const torch::Tensor& Q, - const torch::Tensor& K, - const torch::Tensor& V, - const torch::Tensor& O, - torch::Tensor& dO, - const torch::Tensor& softmax_lse, - torch::Tensor& dK, - torch::Tensor& dV, - float softmax_scale, - bool is_causal, - cudaStream_t stream -) { - using Config = dKVKernelConfig; - - const int B = Q.size(0); - const int H = Q.size(1); - const int M = Q.size(2); - const int N = K.size(2); - - const int grid_x = (N + Config::BLOCK_M - 1) / Config::BLOCK_M; - const dim3 grid(grid_x, 1, B * H); - const dim3 block(Config::THREADS_PER_BLOCK); - const size_t smem = Config::TOTAL_SMEM; - - TORCH_CHECK(smem <= MAX_SMEM_PER_SM, "Shared memory exceeds 96KB for unified dKV: ", smem, " bytes (", smem / 1024, " KB)"); - - auto kernel = is_causal ? - (void*)flash_attention_backward_dkv_kernel : - (void*)flash_attention_backward_dkv_kernel; - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - - if (is_causal) { - flash_attention_backward_dkv_kernel<<>>( + const torch::Tensor& Q, const torch::Tensor& K, const torch::Tensor& V, + const torch::Tensor& O, torch::Tensor& dO, const torch::Tensor& softmax_lse, + torch::Tensor& dK, torch::Tensor& dV, float softmax_scale, bool is_causal, + cudaStream_t stream) { + using Config = dKVKernelConfig; + + const int B = Q.size(0); + const int H = Q.size(1); + const int M = Q.size(2); + const int N = K.size(2); + + const int grid_x = (N + Config::BLOCK_M - 1) / Config::BLOCK_M; + const dim3 grid(grid_x, 1, B * H); + const dim3 block(Config::THREADS_PER_BLOCK); + const size_t smem = Config::TOTAL_SMEM; + + TORCH_CHECK(smem <= MAX_SMEM_PER_SM, + "Shared memory exceeds 96KB for unified dKV: ", smem, " bytes (", + smem / 1024, " KB)"); + + auto kernel = is_causal + ? (void*)flash_attention_backward_dkv_kernel + : (void*)flash_attention_backward_dkv_kernel; + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + if (is_causal) { + flash_attention_backward_dkv_kernel<<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K.data_ptr()), + reinterpret_cast(V.data_ptr()), + reinterpret_cast(O.data_ptr()), + reinterpret_cast<__half*>(dO.data_ptr()), softmax_lse.data_ptr(), + reinterpret_cast<__half*>(dK.data_ptr()), + reinterpret_cast<__half*>(dV.data_ptr()), B, H, M, N, softmax_scale); + } else { + flash_attention_backward_dkv_kernel + <<>>( reinterpret_cast(Q.data_ptr()), reinterpret_cast(K.data_ptr()), reinterpret_cast(V.data_ptr()), @@ -1191,104 +1250,104 @@ void launcher_flash_attention_backward_dkv( reinterpret_cast<__half*>(dO.data_ptr()), softmax_lse.data_ptr(), reinterpret_cast<__half*>(dK.data_ptr()), - reinterpret_cast<__half*>(dV.data_ptr()), - B, H, M, N, softmax_scale - ); - } else { - flash_attention_backward_dkv_kernel<<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K.data_ptr()), - reinterpret_cast(V.data_ptr()), - reinterpret_cast(O.data_ptr()), - reinterpret_cast<__half*>(dO.data_ptr()), - softmax_lse.data_ptr(), - reinterpret_cast<__half*>(dK.data_ptr()), - reinterpret_cast<__half*>(dV.data_ptr()), - B, H, M, N, softmax_scale - ); - } + reinterpret_cast<__half*>(dV.data_ptr()), B, H, M, N, + softmax_scale); + } } std::vector flash_attention_backward( - const at::Tensor& dout, - const at::Tensor& q, - const at::Tensor& k, - const at::Tensor& v, - const at::Tensor& out, - const at::Tensor& softmax_lse, - std::optional& dq_, - std::optional& dk_, - std::optional& dv_, - std::optional& alibi_slopes_, - const float p_dropout, - const float softmax_scale, - const bool is_causal, - int window_size_left, - int window_size_right, - const float softcap, - const bool deterministic, - std::optional gen_, - std::optional& rng_state -) { - - TORCH_CHECK(!alibi_slopes_.has_value(), "alibi_slopes not supported"); - TORCH_CHECK(p_dropout == 0.f, "dropout not supported"); - TORCH_CHECK(window_size_left == -1, "window_size_left not supported"); - TORCH_CHECK(window_size_right == -1 || (is_causal && window_size_right == 0), "window not supported"); - TORCH_CHECK(softcap == 0.f, "softcap not supported"); - TORCH_CHECK(!deterministic, "deterministic mode not supported"); - TORCH_CHECK(!gen_.has_value(), "Generator not supported"); - TORCH_CHECK(!rng_state.has_value() || rng_state->numel() == 0, "rng_state not supported"); - - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - TORCH_CHECK(k.dtype() == torch::kFloat16, "k must be fp16"); - TORCH_CHECK(v.dtype() == torch::kFloat16, "v must be fp16"); - TORCH_CHECK(dout.dtype() == torch::kFloat16, "dout must be fp16"); - TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); - TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, "softmax_lse must be fp32"); - - const auto sizes = q.sizes(); - const int B = sizes[0], H = sizes[1], M = sizes[2], D = sizes[3]; - const int N = k.size(2); - TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, "D must be even, <=256, multiple of 8"); - - at::Tensor dq_fp16 = dq_.has_value() ? dq_.value() : torch::zeros_like(q); - at::Tensor dk_fp16 = dk_.has_value() ? dk_.value() : torch::zeros_like(k); - at::Tensor dv_fp16 = dv_.has_value() ? dv_.value() : torch::zeros_like(v); - - TORCH_CHECK(dq_fp16.dtype() == torch::kFloat16, "dq must be fp16"); - TORCH_CHECK(dk_fp16.dtype() == torch::kFloat16, "dk must be fp16"); - TORCH_CHECK(dv_fp16.dtype() == torch::kFloat16, "dv must be fp16"); - - auto dsoftmax_sum = torch::zeros({B, H, M}, torch::dtype(torch::kFloat32).device(q.device())); - - auto stream = at::cuda::getCurrentCUDAStream().stream(); - auto props = at::cuda::getCurrentDeviceProperties(); - bool sm70 = props->major == 7 && props->minor == 0; - TORCH_CHECK(sm70, "Kernel supports only Volta GPUs."); - - switch (D) { - case 16: - launcher_flash_attention_backward_dq<16>(q, k, v, out, const_cast(dout), softmax_lse, dq_fp16, softmax_scale, is_causal, stream); - launcher_flash_attention_backward_dkv<16>(q, k, v, out, const_cast(dout), softmax_lse, dk_fp16, dv_fp16, softmax_scale, is_causal, stream); - break; - case 32: - launcher_flash_attention_backward_dq<32>(q, k, v, out, const_cast(dout), softmax_lse, dq_fp16, softmax_scale, is_causal, stream); - launcher_flash_attention_backward_dkv<32>(q, k, v, out, const_cast(dout), softmax_lse, dk_fp16, dv_fp16, softmax_scale, is_causal, stream); - break; - case 64: - launcher_flash_attention_backward_dq<64>(q, k, v, out, const_cast(dout), softmax_lse, dq_fp16, softmax_scale, is_causal, stream); - launcher_flash_attention_backward_dkv<64>(q, k, v, out, const_cast(dout), softmax_lse, dk_fp16, dv_fp16, softmax_scale, is_causal, stream); - break; - case 128: - launcher_flash_attention_backward_dq<128>(q, k, v, out, const_cast(dout), softmax_lse, dq_fp16, softmax_scale, is_causal, stream); - launcher_flash_attention_backward_dkv<128>(q, k, v, out, const_cast(dout), softmax_lse, dk_fp16, dv_fp16, softmax_scale, is_causal, stream); - break; - case 256: - launcher_flash_attention_backward_dq<256>(q, k, v, out, const_cast(dout), softmax_lse, dq_fp16, softmax_scale, is_causal, stream); - launcher_flash_attention_backward_dkv<256>(q, k, v, out, const_cast(dout), softmax_lse, dk_fp16, dv_fp16, softmax_scale, is_causal, stream); - break; - default: TORCH_CHECK(false, "Unsupported D: ", D); - } - return {dq_fp16, dk_fp16, dv_fp16, dsoftmax_sum}; + const at::Tensor& dout, const at::Tensor& q, const at::Tensor& k, + const at::Tensor& v, const at::Tensor& out, const at::Tensor& softmax_lse, + std::optional& dq_, std::optional& dk_, + std::optional& dv_, std::optional& alibi_slopes_, + const float p_dropout, const float softmax_scale, const bool is_causal, + int window_size_left, int window_size_right, const float softcap, + const bool deterministic, std::optional gen_, + std::optional& rng_state) { + TORCH_CHECK(!alibi_slopes_.has_value(), "alibi_slopes not supported"); + TORCH_CHECK(p_dropout == 0.f, "dropout not supported"); + TORCH_CHECK(window_size_left == -1, "window_size_left not supported"); + TORCH_CHECK(window_size_right == -1 || (is_causal && window_size_right == 0), + "window not supported"); + TORCH_CHECK(softcap == 0.f, "softcap not supported"); + TORCH_CHECK(!deterministic, "deterministic mode not supported"); + TORCH_CHECK(!gen_.has_value(), "Generator not supported"); + TORCH_CHECK(!rng_state.has_value() || rng_state->numel() == 0, + "rng_state not supported"); + + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + TORCH_CHECK(k.dtype() == torch::kFloat16, "k must be fp16"); + TORCH_CHECK(v.dtype() == torch::kFloat16, "v must be fp16"); + TORCH_CHECK(dout.dtype() == torch::kFloat16, "dout must be fp16"); + TORCH_CHECK(out.dtype() == torch::kFloat16, "out must be fp16"); + TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, + "softmax_lse must be fp32"); + + const auto sizes = q.sizes(); + const int B = sizes[0], H = sizes[1], M = sizes[2], D = sizes[3]; + const int N = k.size(2); + TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, + "D must be even, <=256, multiple of 8"); + + at::Tensor dq_fp16 = dq_.has_value() ? dq_.value() : torch::zeros_like(q); + at::Tensor dk_fp16 = dk_.has_value() ? dk_.value() : torch::zeros_like(k); + at::Tensor dv_fp16 = dv_.has_value() ? dv_.value() : torch::zeros_like(v); + + TORCH_CHECK(dq_fp16.dtype() == torch::kFloat16, "dq must be fp16"); + TORCH_CHECK(dk_fp16.dtype() == torch::kFloat16, "dk must be fp16"); + TORCH_CHECK(dv_fp16.dtype() == torch::kFloat16, "dv must be fp16"); + + auto dsoftmax_sum = + torch::zeros({B, H, M}, torch::dtype(torch::kFloat32).device(q.device())); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + auto props = at::cuda::getCurrentDeviceProperties(); + bool sm70 = props->major == 7 && props->minor == 0; + TORCH_CHECK(sm70, "Kernel supports only Volta GPUs."); + + switch (D) { + case 16: + launcher_flash_attention_backward_dq<16>( + q, k, v, out, const_cast(dout), softmax_lse, dq_fp16, + softmax_scale, is_causal, stream); + launcher_flash_attention_backward_dkv<16>( + q, k, v, out, const_cast(dout), softmax_lse, dk_fp16, + dv_fp16, softmax_scale, is_causal, stream); + break; + case 32: + launcher_flash_attention_backward_dq<32>( + q, k, v, out, const_cast(dout), softmax_lse, dq_fp16, + softmax_scale, is_causal, stream); + launcher_flash_attention_backward_dkv<32>( + q, k, v, out, const_cast(dout), softmax_lse, dk_fp16, + dv_fp16, softmax_scale, is_causal, stream); + break; + case 64: + launcher_flash_attention_backward_dq<64>( + q, k, v, out, const_cast(dout), softmax_lse, dq_fp16, + softmax_scale, is_causal, stream); + launcher_flash_attention_backward_dkv<64>( + q, k, v, out, const_cast(dout), softmax_lse, dk_fp16, + dv_fp16, softmax_scale, is_causal, stream); + break; + case 128: + launcher_flash_attention_backward_dq<128>( + q, k, v, out, const_cast(dout), softmax_lse, dq_fp16, + softmax_scale, is_causal, stream); + launcher_flash_attention_backward_dkv<128>( + q, k, v, out, const_cast(dout), softmax_lse, dk_fp16, + dv_fp16, softmax_scale, is_causal, stream); + break; + case 256: + launcher_flash_attention_backward_dq<256>( + q, k, v, out, const_cast(dout), softmax_lse, dq_fp16, + softmax_scale, is_causal, stream); + launcher_flash_attention_backward_dkv<256>( + q, k, v, out, const_cast(dout), softmax_lse, dk_fp16, + dv_fp16, softmax_scale, is_causal, stream); + break; + default: + TORCH_CHECK(false, "Unsupported D: ", D); + } + return {dq_fp16, dk_fp16, dv_fp16, dsoftmax_sum}; } diff --git a/flash-attention-v100/kernel/fused_mha_forward.cu b/flash-attention-v100/kernel/fused_mha_forward.cu index 1bc0c75448..fc9cfb4c7a 100644 --- a/flash-attention-v100/kernel/fused_mha_forward.cu +++ b/flash-attention-v100/kernel/fused_mha_forward.cu @@ -16,1133 +16,1109 @@ using namespace nvcuda::wmma; #define WMMA_N 16 #define WMMA_K 16 -#define MAX_THREADS_PER_WARP 32 -#define MAX_THREADS_PER_SM 2048 -#define MAX_THREAD_BLOCK_SIZE 1024 +#define MAX_THREADS_PER_WARP 32 +#define MAX_THREADS_PER_SM 2048 +#define MAX_THREAD_BLOCK_SIZE 1024 #define MAX_THREAD_BLOCK_PER_SM 32 -#define MAX_WARPS_PER_SM 64 -#define MAX_SM_PER_GPU 80 -#define MAX_SMEM_PER_SM 98304 +#define MAX_WARPS_PER_SM 64 +#define MAX_SM_PER_GPU 80 +#define MAX_SMEM_PER_SM 98304 -#define WARP_ALLOC_GROUP 4 +#define WARP_ALLOC_GROUP 4 -#define MAX_REG_PER_UNIT 256 -#define MAX_REG_PER_THREAD 255 -#define MAX_REG_PER_BLOCK 65536 -#define MAX_REG_BUFFER 65536 +#define MAX_REG_PER_UNIT 256 +#define MAX_REG_PER_THREAD 255 +#define MAX_REG_PER_BLOCK 65536 +#define MAX_REG_BUFFER 65536 -#define BLOCK_M_16 16 -#define BLOCK_N_16 512 -#define WARPS_16 16 +#define BLOCK_M_16 16 +#define BLOCK_N_16 512 +#define WARPS_16 16 -#define BLOCK_M_32 32 -#define BLOCK_N_32 256 -#define WARPS_32 16 +#define BLOCK_M_32 32 +#define BLOCK_N_32 256 +#define WARPS_32 16 -#define BLOCK_M_64 64 -#define BLOCK_N_64 128 -#define WARPS_64 16 +#define BLOCK_M_64 64 +#define BLOCK_N_64 128 +#define WARPS_64 16 #define BLOCK_M_128 32 #define BLOCK_N_128 176 -#define WARPS_128 16 +#define WARPS_128 16 #define BLOCK_M_256 32 #define BLOCK_N_256 64 -#define WARPS_256 16 +#define WARPS_256 16 #define BLOCK_M_256_LOW_SMEM 16 #define BLOCK_N_256_LOW_SMEM 32 -#define WARPS_256_LOW_SMEM 16 +#define WARPS_256_LOW_SMEM 16 inline bool env_flag_enabled(const char* name) { - const char* raw = std::getenv(name); - return raw != nullptr && std::strcmp(raw, "0") != 0; + const char* raw = std::getenv(name); + return raw != nullptr && std::strcmp(raw, "0") != 0; } inline bool env_flag_enabled_default_on(const char* name) { - const char* raw = std::getenv(name); - return raw == nullptr || std::strcmp(raw, "0") != 0; + const char* raw = std::getenv(name); + return raw == nullptr || std::strcmp(raw, "0") != 0; } -template +template struct KernelConfig { - static constexpr int BLOCK_M = (D == 16) ? BLOCK_M_16 : (D == 32) ? BLOCK_M_32 : (D == 64) ? BLOCK_M_64 : (D == 128) ? BLOCK_M_128 : (LOW_SMEM ? BLOCK_M_256_LOW_SMEM : BLOCK_M_256); - static constexpr int BLOCK_N = (D == 16) ? BLOCK_N_16 : (D == 32) ? BLOCK_N_32 : (D == 64) ? BLOCK_N_64 : (D == 128) ? BLOCK_N_128 : (LOW_SMEM ? BLOCK_N_256_LOW_SMEM : BLOCK_N_256); - static constexpr int WARPS_PER_BLOCK = (D == 16) ? WARPS_16 : (D == 32) ? WARPS_32 : (D == 64) ? WARPS_64 : (D == 128) ? WARPS_128 : (LOW_SMEM ? WARPS_256_LOW_SMEM : WARPS_256); - - static constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * MAX_THREADS_PER_WARP; - static constexpr int THREADS_PER_ROW = THREADS_PER_BLOCK / BLOCK_M; - static constexpr int PAD = (8 - (D % 32) + 32) % 32; - static constexpr int Q_STRIDE = D + PAD; - static constexpr int KV_STRIDE = D + PAD; - static constexpr int S_STRIDE = BLOCK_N + PAD; - static constexpr int P_SUB_TILE = 32; - static constexpr int P_STRIDE = P_SUB_TILE + PAD; - static constexpr int P_STRICT_ELEMENTS = (D == 256) ? BLOCK_M * P_STRIDE : 1; - static constexpr int O_STRIDE = D + PAD; - static constexpr int PER_UINT4 = 8; - - struct alignas(128) SmemLayout { - alignas(16) __half q [BLOCK_M * Q_STRIDE]; + static constexpr int BLOCK_M = + (D == 16) ? BLOCK_M_16 + : (D == 32) ? BLOCK_M_32 + : (D == 64) ? BLOCK_M_64 + : (D == 128) ? BLOCK_M_128 + : (LOW_SMEM ? BLOCK_M_256_LOW_SMEM : BLOCK_M_256); + static constexpr int BLOCK_N = + (D == 16) ? BLOCK_N_16 + : (D == 32) ? BLOCK_N_32 + : (D == 64) ? BLOCK_N_64 + : (D == 128) ? BLOCK_N_128 + : (LOW_SMEM ? BLOCK_N_256_LOW_SMEM : BLOCK_N_256); + static constexpr int WARPS_PER_BLOCK = + (D == 16) ? WARPS_16 + : (D == 32) ? WARPS_32 + : (D == 64) ? WARPS_64 + : (D == 128) ? WARPS_128 + : (LOW_SMEM ? WARPS_256_LOW_SMEM : WARPS_256); + + static constexpr int THREADS_PER_BLOCK = + WARPS_PER_BLOCK * MAX_THREADS_PER_WARP; + static constexpr int THREADS_PER_ROW = THREADS_PER_BLOCK / BLOCK_M; + static constexpr int PAD = (8 - (D % 32) + 32) % 32; + static constexpr int Q_STRIDE = D + PAD; + static constexpr int KV_STRIDE = D + PAD; + static constexpr int S_STRIDE = BLOCK_N + PAD; + static constexpr int P_SUB_TILE = 32; + static constexpr int P_STRIDE = P_SUB_TILE + PAD; + static constexpr int P_STRICT_ELEMENTS = (D == 256) ? BLOCK_M * P_STRIDE : 1; + static constexpr int O_STRIDE = D + PAD; + static constexpr int PER_UINT4 = 8; + + struct alignas(128) SmemLayout { + alignas(16) __half q[BLOCK_M * Q_STRIDE]; union { - alignas(16) __half k [BLOCK_N * KV_STRIDE]; - alignas(16) __half v [BLOCK_N * KV_STRIDE]; + alignas(16) __half k[BLOCK_N * KV_STRIDE]; + alignas(16) __half v[BLOCK_N * KV_STRIDE]; } reuse_kv; union { - alignas(16) float s [BLOCK_M * S_STRIDE]; - alignas(16) __half p [BLOCK_M * S_STRIDE]; + alignas(16) float s[BLOCK_M * S_STRIDE]; + alignas(16) __half p[BLOCK_M * S_STRIDE]; } reuse_sp; - alignas(16) __half p_strict[P_STRICT_ELEMENTS]; - alignas(16) float o [BLOCK_M * O_STRIDE]; - alignas(16) float row_max[BLOCK_M]; - alignas(16) float row_sum[BLOCK_M]; - }; - - static constexpr size_t TOTAL_SMEM = ((sizeof(SmemLayout) + 127) & ~size_t(127)); + alignas(16) __half p_strict[P_STRICT_ELEMENTS]; + alignas(16) float o[BLOCK_M * O_STRIDE]; + alignas(16) float row_max[BLOCK_M]; + alignas(16) float row_sum[BLOCK_M]; + }; + + static constexpr size_t TOTAL_SMEM = + ((sizeof(SmemLayout) + 127) & ~size_t(127)); }; -template +template __device__ __forceinline__ void init_smem(char* smem_raw) { - constexpr int N_U4 = Config::TOTAL_SMEM / 16; - constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; - const int tid = threadIdx.x; - - uint32_t addr = static_cast(__cvta_generic_to_shared(smem_raw)); - #pragma unroll 1 - for (int i = tid; i < N_U4; i += THREADS_PER_BLOCK) { - asm volatile("st.shared.v4.u32 [%0], {%1,%1,%1,%1};" - :: "r"(addr + (i << 4)), "r"(0) : "memory"); - } - __syncthreads(); + constexpr int N_U4 = Config::TOTAL_SMEM / 16; + constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; + const int tid = threadIdx.x; + + uint32_t addr = static_cast(__cvta_generic_to_shared(smem_raw)); +#pragma unroll 1 + for (int i = tid; i < N_U4; i += THREADS_PER_BLOCK) { + asm volatile("st.shared.v4.u32 [%0], {%1,%1,%1,%1};" ::"r"(addr + (i << 4)), + "r"(0) + : "memory"); + } + __syncthreads(); } -template +template __device__ __forceinline__ void compute_qk_scores_scalar_fp32( - const __half* __restrict__ sQ, - const __half* __restrict__ sK, - float* __restrict__ sS, - const int valid_q_rows, - const int valid_k_rows, - const int start_row, - const int start_col, - const int causal_q_offset, - const float score_scale, - const int window_size_left, - const int window_size_right -) { - using Config = KernelConfig; - constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; - constexpr int Q_STRIDE = Config::Q_STRIDE; - constexpr int KV_STRIDE = Config::KV_STRIDE; - constexpr int S_STRIDE = Config::S_STRIDE; - const float NEG_INF = -1e30f; - const int tid = threadIdx.x; - const int total_scores = valid_q_rows * valid_k_rows; - - for (int idx = tid; idx < total_scores; idx += THREADS_PER_BLOCK) { - const int row = idx / valid_k_rows; - const int col = idx - row * valid_k_rows; - const int global_m = start_row + row; - const int global_n = start_col + col; - const int global_q_pos = global_m + causal_q_offset; - - bool is_valid = true; - if constexpr (IS_CAUSAL) { - is_valid = global_n <= global_q_pos; - } - if (window_size_left >= 0) { - is_valid = is_valid && global_n >= global_q_pos - window_size_left; - } - if (window_size_right >= 0) { - is_valid = is_valid && global_n <= global_q_pos + window_size_right; - } - - float acc = 0.0f; - if (is_valid) { - #pragma unroll 8 - for (int d = 0; d < D; d += 2) { - const __half2 q_h2 = *reinterpret_cast( - sQ + row * Q_STRIDE + d); - const __half2 k_h2 = *reinterpret_cast( - sK + col * KV_STRIDE + d); - const float2 q_f2 = __half22float2(q_h2); - const float2 k_f2 = __half22float2(k_h2); - acc = fmaf(q_f2.x, k_f2.x, acc); - acc = fmaf(q_f2.y, k_f2.y, acc); - } - } - sS[row * S_STRIDE + col] = is_valid ? acc * score_scale : NEG_INF; - } -} - -template -__global__ void __launch_bounds__(KernelConfig::THREADS_PER_BLOCK, 2) -flash_attention_forward_kernel( - const __half* __restrict__ Q, - const __half* __restrict__ K, - const __half* __restrict__ V, - __half* __restrict__ Out, - float* __restrict__ softmax_lse, - const int B, - const int H, - const int H_KV, - const int M, - const int N, - const float softmax_scale, - const bool scalar_pv, - const int window_size_left, - const int window_size_right -) { - using Config = KernelConfig; - constexpr int BLOCK_M = Config::BLOCK_M; - constexpr int BLOCK_N = Config::BLOCK_N; - constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; - constexpr int THREADS_PER_ROW = Config::THREADS_PER_ROW; - constexpr int WARPS_PER_BLOCK = Config::WARPS_PER_BLOCK; - constexpr int Q_STRIDE = Config::Q_STRIDE; - constexpr int KV_STRIDE = Config::KV_STRIDE; - constexpr int S_STRIDE = Config::S_STRIDE; - constexpr int P_SUB_TILE = Config::P_SUB_TILE; - constexpr int P_STRIDE = Config::P_STRIDE; - constexpr int O_STRIDE = Config::O_STRIDE; - constexpr int PER_UINT4 = Config::PER_UINT4; - const float NEG_INF = -1e30f; - - const int batch_head_id = blockIdx.z; - if (batch_head_id >= B * H) return; - const int batch_id = batch_head_id / H; - const int q_head_id = batch_head_id % H; - const int kv_group_size = H / H_KV; - const int kv_head_id = q_head_id / kv_group_size; - - const int block_m = blockIdx.x; - const int start_row = block_m * BLOCK_M; - if (start_row >= M) return; - - int num_n_tiles = (N + BLOCK_N - 1) / BLOCK_N; - const int valid_q_rows = min(BLOCK_M, M - start_row); - const int causal_q_offset = max(N - M, 0); - - int max_key_pos = N - 1; + const __half* __restrict__ sQ, const __half* __restrict__ sK, + float* __restrict__ sS, const int valid_q_rows, const int valid_k_rows, + const int start_row, const int start_col, const int causal_q_offset, + const float score_scale, const int window_size_left, + const int window_size_right) { + using Config = KernelConfig; + constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; + constexpr int Q_STRIDE = Config::Q_STRIDE; + constexpr int KV_STRIDE = Config::KV_STRIDE; + constexpr int S_STRIDE = Config::S_STRIDE; + const float NEG_INF = -1e30f; + const int tid = threadIdx.x; + const int total_scores = valid_q_rows * valid_k_rows; + + for (int idx = tid; idx < total_scores; idx += THREADS_PER_BLOCK) { + const int row = idx / valid_k_rows; + const int col = idx - row * valid_k_rows; + const int global_m = start_row + row; + const int global_n = start_col + col; + const int global_q_pos = global_m + causal_q_offset; + + bool is_valid = true; if constexpr (IS_CAUSAL) { - max_key_pos = min(max_key_pos, - causal_q_offset + start_row + valid_q_rows - 1); + is_valid = global_n <= global_q_pos; + } + if (window_size_left >= 0) { + is_valid = is_valid && global_n >= global_q_pos - window_size_left; } if (window_size_right >= 0) { - max_key_pos = min(max_key_pos, - causal_q_offset + start_row + valid_q_rows - 1 - + window_size_right); + is_valid = is_valid && global_n <= global_q_pos + window_size_right; } - if (max_key_pos < 0) { - num_n_tiles = 0; - } else { - num_n_tiles = min(num_n_tiles, (max_key_pos + BLOCK_N) / BLOCK_N); + + float acc = 0.0f; + if (is_valid) { +#pragma unroll 8 + for (int d = 0; d < D; d += 2) { + const __half2 q_h2 = + *reinterpret_cast(sQ + row * Q_STRIDE + d); + const __half2 k_h2 = + *reinterpret_cast(sK + col * KV_STRIDE + d); + const float2 q_f2 = __half22float2(q_h2); + const float2 k_f2 = __half22float2(k_h2); + acc = fmaf(q_f2.x, k_f2.x, acc); + acc = fmaf(q_f2.y, k_f2.y, acc); + } } - const int min_key_pos = window_size_left >= 0 - ? max(0, causal_q_offset + start_row - - window_size_left) - : 0; - - const int tid = threadIdx.x; - const int warp_id = tid / MAX_THREADS_PER_WARP; - const int lane_id = tid % MAX_THREADS_PER_WARP; - - const size_t q_head_linear = (size_t)batch_id * H + q_head_id; - const size_t kv_head_linear = (size_t)batch_id * H_KV + kv_head_id; - const __half* q_ptr = Q + q_head_linear * M * D + start_row * D; - const __half* k_ptr = K + kv_head_linear * N * D; - const __half* v_ptr = V + kv_head_linear * N * D; - __half* out_ptr = Out + q_head_linear * M * D + start_row * D; - float* softmax_lse_ptr = softmax_lse + q_head_linear * M + start_row; - - extern __shared__ char smem_raw[]; - init_smem(smem_raw); - auto& smem = *reinterpret_cast(smem_raw); - - __half* sQ = smem.q; - __half* sK = smem.reuse_kv.k; - __half* sV = smem.reuse_kv.v; - float* sS = smem.reuse_sp.s; - __half* sP = (D == 256) ? smem.p_strict : smem.reuse_sp.p; - const int p_stride = (D == 256) ? P_STRIDE : S_STRIDE; - const int p_tile_capacity = (D == 256) ? P_SUB_TILE : BLOCK_N; - float* sO = smem.o; - float* sRowMax = smem.row_max; - float* sRowSum = smem.row_sum; - constexpr float RCP_LN2 = 1.4426950408889634f; - constexpr float LN2 = 0.6931471805599453f; - constexpr float EXP2_CLAMP = -80.0f * RCP_LN2; - const float softmax_scale_log2 = softmax_scale * RCP_LN2; - - const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; - const int q_stride_uint4 = (Q_STRIDE + PER_UINT4 - 1) / PER_UINT4; - const int kv_stride_uint4 = (KV_STRIDE + PER_UINT4 - 1) / PER_UINT4; - - if (tid < BLOCK_M) { - sRowMax[tid] = NEG_INF; + sS[row * S_STRIDE + col] = is_valid ? acc * score_scale : NEG_INF; + } +} + +template +__global__ void __launch_bounds__(KernelConfig::THREADS_PER_BLOCK, + 2) + flash_attention_forward_kernel( + const __half* __restrict__ Q, const __half* __restrict__ K, + const __half* __restrict__ V, __half* __restrict__ Out, + float* __restrict__ softmax_lse, const int B, const int H, + const int H_KV, const int M, const int N, const float softmax_scale, + const bool scalar_pv, const int window_size_left, + const int window_size_right) { + using Config = KernelConfig; + constexpr int BLOCK_M = Config::BLOCK_M; + constexpr int BLOCK_N = Config::BLOCK_N; + constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; + constexpr int THREADS_PER_ROW = Config::THREADS_PER_ROW; + constexpr int WARPS_PER_BLOCK = Config::WARPS_PER_BLOCK; + constexpr int Q_STRIDE = Config::Q_STRIDE; + constexpr int KV_STRIDE = Config::KV_STRIDE; + constexpr int S_STRIDE = Config::S_STRIDE; + constexpr int P_SUB_TILE = Config::P_SUB_TILE; + constexpr int P_STRIDE = Config::P_STRIDE; + constexpr int O_STRIDE = Config::O_STRIDE; + constexpr int PER_UINT4 = Config::PER_UINT4; + const float NEG_INF = -1e30f; + + const int batch_head_id = blockIdx.z; + if (batch_head_id >= B * H) return; + const int batch_id = batch_head_id / H; + const int q_head_id = batch_head_id % H; + const int kv_group_size = H / H_KV; + const int kv_head_id = q_head_id / kv_group_size; + + const int block_m = blockIdx.x; + const int start_row = block_m * BLOCK_M; + if (start_row >= M) return; + + int num_n_tiles = (N + BLOCK_N - 1) / BLOCK_N; + const int valid_q_rows = min(BLOCK_M, M - start_row); + const int causal_q_offset = max(N - M, 0); + + int max_key_pos = N - 1; + if constexpr (IS_CAUSAL) { + max_key_pos = + min(max_key_pos, causal_q_offset + start_row + valid_q_rows - 1); + } + if (window_size_right >= 0) { + max_key_pos = min(max_key_pos, causal_q_offset + start_row + valid_q_rows - + 1 + window_size_right); + } + if (max_key_pos < 0) { + num_n_tiles = 0; + } else { + num_n_tiles = min(num_n_tiles, (max_key_pos + BLOCK_N) / BLOCK_N); + } + const int min_key_pos = + window_size_left >= 0 + ? max(0, causal_q_offset + start_row - window_size_left) + : 0; + + const int tid = threadIdx.x; + const int warp_id = tid / MAX_THREADS_PER_WARP; + const int lane_id = tid % MAX_THREADS_PER_WARP; + + const size_t q_head_linear = (size_t)batch_id * H + q_head_id; + const size_t kv_head_linear = (size_t)batch_id * H_KV + kv_head_id; + const __half* q_ptr = Q + q_head_linear * M * D + start_row * D; + const __half* k_ptr = K + kv_head_linear * N * D; + const __half* v_ptr = V + kv_head_linear * N * D; + __half* out_ptr = Out + q_head_linear * M * D + start_row * D; + float* softmax_lse_ptr = softmax_lse + q_head_linear * M + start_row; + + extern __shared__ char smem_raw[]; + init_smem(smem_raw); + auto& smem = *reinterpret_cast(smem_raw); + + __half* sQ = smem.q; + __half* sK = smem.reuse_kv.k; + __half* sV = smem.reuse_kv.v; + float* sS = smem.reuse_sp.s; + __half* sP = (D == 256) ? smem.p_strict : smem.reuse_sp.p; + const int p_stride = (D == 256) ? P_STRIDE : S_STRIDE; + const int p_tile_capacity = (D == 256) ? P_SUB_TILE : BLOCK_N; + float* sO = smem.o; + float* sRowMax = smem.row_max; + float* sRowSum = smem.row_sum; + constexpr float RCP_LN2 = 1.4426950408889634f; + constexpr float LN2 = 0.6931471805599453f; + constexpr float EXP2_CLAMP = -80.0f * RCP_LN2; + const float softmax_scale_log2 = softmax_scale * RCP_LN2; + + const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; + const int q_stride_uint4 = (Q_STRIDE + PER_UINT4 - 1) / PER_UINT4; + const int kv_stride_uint4 = (KV_STRIDE + PER_UINT4 - 1) / PER_UINT4; + + if (tid < BLOCK_M) { + sRowMax[tid] = NEG_INF; + } + + const uint4* q_vec = reinterpret_cast(q_ptr); + uint4* sQ_vec = reinterpret_cast(sQ); + +#pragma unroll 4 + for (int idx = tid; idx < (valid_q_rows * d_stride_uint4); + idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + uint4 q_val = make_uint4(0, 0, 0, 0); + if (row < valid_q_rows && vec_col < d_stride_uint4) { + q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); } + sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + } + __syncthreads(); - const uint4* q_vec = reinterpret_cast(q_ptr); - uint4* sQ_vec = reinterpret_cast(sQ); + for (int block_n = 0; block_n < num_n_tiles; ++block_n) { + const int start_col = block_n * BLOCK_N; + if (start_col >= N) break; + const int valid_k_rows = min(BLOCK_N, N - start_col); - #pragma unroll 4 - for (int idx = tid; idx < (valid_q_rows * d_stride_uint4); idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - uint4 q_val = make_uint4(0, 0, 0, 0); - if (row < valid_q_rows && vec_col < d_stride_uint4) { - q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); - } - sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + if (start_col + valid_k_rows <= min_key_pos) { + continue; + } + + const uint4* k_vec = reinterpret_cast(k_ptr + start_col * D); + uint4* sK_vec = reinterpret_cast(sK); + +#pragma unroll 2 + for (int idx = tid; idx < (valid_k_rows * d_stride_uint4); + idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + + uint4 k_val = make_uint4(0, 0, 0, 0); + if (row < valid_k_rows && vec_col < d_stride_uint4) { + k_val = __ldg(&k_vec[row * d_stride_uint4 + vec_col]); + } + sK_vec[row * kv_stride_uint4 + vec_col] = k_val; } __syncthreads(); - for (int block_n = 0; block_n < num_n_tiles; ++block_n) { - const int start_col = block_n * BLOCK_N; - if (start_col >= N) break; - const int valid_k_rows = min(BLOCK_N, N - start_col); + if constexpr (D == 256 && !D256_WMMA_QK) { + compute_qk_scores_scalar_fp32( + sQ, sK, sS, valid_q_rows, valid_k_rows, start_row, start_col, + causal_q_offset, softmax_scale, window_size_left, window_size_right); + } else { + const int num_tiles_m_qk = (BLOCK_M + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_qk = (BLOCK_N + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_qk = (D + WMMA_K - 1) / WMMA_K; + const int total_tiles_qk = num_tiles_m_qk * num_tiles_n_qk; + const int tiles_per_warp_qk = + (total_tiles_qk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + const unsigned row_causal = (lane_id & 0b1) + ((lane_id >> 2) & 0b1) * 8 + + ((lane_id >> 4) & 0b1) * 4; + const unsigned col_causal = + ((lane_id >> 1) & 0b1) * 2 + ((lane_id >> 3) & 0b1) * 8; + + for (int tile_idx = 0; tile_idx < tiles_per_warp_qk; ++tile_idx) { + const int global_tile_idx = warp_id * tiles_per_warp_qk + tile_idx; + if (global_tile_idx >= total_tiles_qk) break; + + const int tile_m_idx = global_tile_idx / num_tiles_n_qk; + const int tile_n_idx = global_tile_idx % num_tiles_n_qk; + + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; + + if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + fill_fragment(acc_frag, 0.0f); + +#pragma unroll + for (int k_tile = 0; k_tile < num_tiles_k_qk; ++k_tile) { + const int k_offset = k_tile * WMMA_K; + if (k_offset >= D) break; + + load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, Q_STRIDE); + load_matrix_sync(b_frag, sK + tile_n * KV_STRIDE + k_offset, + KV_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } - if (start_col + valid_k_rows <= min_key_pos) { - continue; +#pragma unroll + for (int i = 0; i < acc_frag.num_elements; ++i) { + const unsigned col = col_causal + (i & 0b1) + ((i >> 2) & 0b1) * 4; + const unsigned row = row_causal + ((i >> 1) & 0b1) * 2; + + const int global_m = start_row + tile_m + row; + const int global_n = start_col + tile_n + col; + const int global_q_pos = global_m + causal_q_offset; + + const bool is_valid = (global_m < start_row + valid_q_rows) && + (global_n < start_col + valid_k_rows); + bool is_causal_valid = true; + if constexpr (IS_CAUSAL) { + is_causal_valid = global_n <= global_q_pos; + } + bool is_window_valid = true; + if (window_size_left >= 0) { + is_window_valid = + is_window_valid && global_n >= global_q_pos - window_size_left; + } + if (window_size_right >= 0) { + is_window_valid = + is_window_valid && global_n <= global_q_pos + window_size_right; + } + + const float qk_scale = + (D == 256) ? softmax_scale : softmax_scale_log2; + acc_frag.x[i] = (is_valid && is_causal_valid && is_window_valid) + ? acc_frag.x[i] * qk_scale + : NEG_INF; } + store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, + mem_row_major); + } + } + __syncthreads(); - const uint4* k_vec = reinterpret_cast(k_ptr + start_col * D); - uint4* sK_vec = reinterpret_cast(sK); + const uint4* v_vec = reinterpret_cast(v_ptr + start_col * D); + uint4* sV_vec = reinterpret_cast(sV); - #pragma unroll 2 - for (int idx = tid; idx < (valid_k_rows * d_stride_uint4); idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; +#pragma unroll 2 + for (int idx = tid; idx < (valid_k_rows * d_stride_uint4); + idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + + uint4 v_val = make_uint4(0, 0, 0, 0); + if (row < valid_k_rows && vec_col < d_stride_uint4) { + v_val = __ldg(&v_vec[row * d_stride_uint4 + vec_col]); + } + sV_vec[row * kv_stride_uint4 + vec_col] = v_val; + } + __syncthreads(); - uint4 k_val = make_uint4(0, 0, 0, 0); - if (row < valid_k_rows && vec_col < d_stride_uint4) { - k_val = __ldg(&k_vec[row * d_stride_uint4 + vec_col]); - } - sK_vec[row * kv_stride_uint4 + vec_col] = k_val; + const int num_tiles_m_pv = (BLOCK_M + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_pv = (D + WMMA_N - 1) / WMMA_N; + const int total_tiles_pv = num_tiles_m_pv * num_tiles_n_pv; + const int tiles_per_warp_pv = + (total_tiles_pv + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + const int softmax_sub_tile = (D == 256) ? P_SUB_TILE : p_tile_capacity; + + for (int sub_start = 0; sub_start < valid_k_rows; + sub_start += softmax_sub_tile) { + const int sub_valid_k_rows = + min(softmax_sub_tile, valid_k_rows - sub_start); + + if (tid < valid_q_rows * THREADS_PER_ROW) { + const int row = tid / THREADS_PER_ROW; + const int thread_in_row = tid % THREADS_PER_ROW; + const unsigned mask = + (valid_q_rows == BLOCK_M) ? 0xFFFFFFFFU : __activemask(); + const int row_leader = __ffs(mask) - 1; + + float* sS_row_f = sS + row * S_STRIDE + sub_start; + __half* sP_row_h = sP + row * p_stride; + + const int vec_cols = sub_valid_k_rows >> 2; + const int vecs_per_thread = + (vec_cols + THREADS_PER_ROW - 1) / THREADS_PER_ROW; + const int tail_start = vec_cols << 2; + + float thread_max = NEG_INF; + float4* sS_vec4 = reinterpret_cast(sS_row_f); + +#pragma unroll 4 + for (int j = 0; j < vecs_per_thread; ++j) { + int vc = thread_in_row + j * THREADS_PER_ROW; + if (vc < vec_cols) { + float4 v4 = sS_vec4[vc]; + thread_max = + fmaxf(thread_max, fmaxf(fmaxf(v4.x, v4.y), fmaxf(v4.z, v4.w))); + } } - __syncthreads(); - if constexpr (D == 256 && !D256_WMMA_QK) { - compute_qk_scores_scalar_fp32( - sQ, sK, sS, valid_q_rows, valid_k_rows, start_row, start_col, - causal_q_offset, softmax_scale, window_size_left, - window_size_right); - } else { - const int num_tiles_m_qk = (BLOCK_M + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_qk = (BLOCK_N + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_qk = (D + WMMA_K - 1) / WMMA_K; - const int total_tiles_qk = num_tiles_m_qk * num_tiles_n_qk; - const int tiles_per_warp_qk = (total_tiles_qk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - const unsigned row_causal = (lane_id & 0b1) + ((lane_id >> 2) & 0b1) * 8 + ((lane_id >> 4) & 0b1) * 4; - const unsigned col_causal = ((lane_id >> 1) & 0b1) * 2 + ((lane_id >> 3) & 0b1) * 8; - - for (int tile_idx = 0; tile_idx < tiles_per_warp_qk; ++tile_idx) { - const int global_tile_idx = warp_id * tiles_per_warp_qk + tile_idx; - if (global_tile_idx >= total_tiles_qk) break; - - const int tile_m_idx = global_tile_idx / num_tiles_n_qk; - const int tile_n_idx = global_tile_idx % num_tiles_n_qk; - - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; - - if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - fill_fragment(acc_frag, 0.0f); - - #pragma unroll - for (int k_tile = 0; k_tile < num_tiles_k_qk; ++k_tile) { - const int k_offset = k_tile * WMMA_K; - if (k_offset >= D) break; - - load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, Q_STRIDE); - load_matrix_sync(b_frag, sK + tile_n * KV_STRIDE + k_offset, KV_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - - #pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - const unsigned col = col_causal + (i & 0b1) + ((i >> 2) & 0b1) * 4; - const unsigned row = row_causal + ((i >> 1) & 0b1) * 2; - - const int global_m = start_row + tile_m + row; - const int global_n = start_col + tile_n + col; - const int global_q_pos = global_m + causal_q_offset; - - const bool is_valid = (global_m < start_row + valid_q_rows) && - (global_n < start_col + valid_k_rows); - bool is_causal_valid = true; - if constexpr (IS_CAUSAL) { - is_causal_valid = global_n <= global_q_pos; - } - bool is_window_valid = true; - if (window_size_left >= 0) { - is_window_valid = is_window_valid && - global_n >= global_q_pos - window_size_left; - } - if (window_size_right >= 0) { - is_window_valid = is_window_valid && - global_n <= global_q_pos + window_size_right; - } - - const float qk_scale = - (D == 256) ? softmax_scale : softmax_scale_log2; - acc_frag.x[i] = - (is_valid && is_causal_valid && is_window_valid) - ? acc_frag.x[i] * qk_scale - : NEG_INF; - } - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, mem_row_major); - } +#pragma unroll 4 + for (int c = tail_start + thread_in_row; c < sub_valid_k_rows; + c += THREADS_PER_ROW) { + thread_max = fmaxf(thread_max, sS_row_f[c]); } - __syncthreads(); - const uint4* v_vec = reinterpret_cast(v_ptr + start_col * D); - uint4* sV_vec = reinterpret_cast(sV); +#pragma unroll + for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) { + thread_max = fmaxf(thread_max, __shfl_down_sync(mask, thread_max, o, + THREADS_PER_ROW)); + } - #pragma unroll 2 - for (int idx = tid; idx < (valid_k_rows * d_stride_uint4); idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; + const float row_max = + __shfl_sync(mask, thread_max, row_leader, THREADS_PER_ROW); + const float old_max = sRowMax[row]; + const float new_max = fmaxf(old_max, row_max); + const float exp_diff = (D == 256) + ? expf(fmaxf(old_max - new_max, -80.0f)) + : exp2f(old_max - new_max); + + float thread_sum = 0.0f; + __half2 half_buffer[20]; + int vc_base = thread_in_row; + int h2_idx = 0; + int tail_col = -1; + __half tail_value = __float2half(0.f); + +#pragma unroll 4 + for (int j = 0; j < vecs_per_thread; ++j, vc_base += THREADS_PER_ROW) { + if (vc_base < vec_cols) { + float4 v4 = sS_vec4[vc_base]; + + float e0 = (D == 256) ? expf(fmaxf(v4.x - new_max, -80.0f)) + : exp2f(fmaxf(v4.x - new_max, EXP2_CLAMP)); + float e1 = (D == 256) ? expf(fmaxf(v4.y - new_max, -80.0f)) + : exp2f(fmaxf(v4.y - new_max, EXP2_CLAMP)); + float e2 = (D == 256) ? expf(fmaxf(v4.z - new_max, -80.0f)) + : exp2f(fmaxf(v4.z - new_max, EXP2_CLAMP)); + float e3 = (D == 256) ? expf(fmaxf(v4.w - new_max, -80.0f)) + : exp2f(fmaxf(v4.w - new_max, EXP2_CLAMP)); + + thread_sum += (e0 + e1) + (e2 + e3); + + half_buffer[h2_idx++] = __float22half2_rn(make_float2(e0, e1)); + half_buffer[h2_idx++] = __float22half2_rn(make_float2(e2, e3)); + } + } - uint4 v_val = make_uint4(0, 0, 0, 0); - if (row < valid_k_rows && vec_col < d_stride_uint4) { - v_val = __ldg(&v_vec[row * d_stride_uint4 + vec_col]); - } - sV_vec[row * kv_stride_uint4 + vec_col] = v_val; +#pragma unroll 4 + for (int c = tail_start + thread_in_row; c < sub_valid_k_rows; + c += THREADS_PER_ROW) { + float v = sS_row_f[c]; + float e = (D == 256) ? expf(fmaxf(v - new_max, -80.0f)) + : exp2f(fmaxf(v - new_max, EXP2_CLAMP)); + thread_sum += e; + tail_col = c; + tail_value = __float2half_rn(e); } - __syncthreads(); - const int num_tiles_m_pv = (BLOCK_M + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_pv = (D + WMMA_N - 1) / WMMA_N; - const int total_tiles_pv = num_tiles_m_pv * num_tiles_n_pv; - const int tiles_per_warp_pv = (total_tiles_pv + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - const int softmax_sub_tile = (D == 256) ? P_SUB_TILE : p_tile_capacity; - - for (int sub_start = 0; sub_start < valid_k_rows; - sub_start += softmax_sub_tile) { - const int sub_valid_k_rows = min(softmax_sub_tile, - valid_k_rows - sub_start); - - if (tid < valid_q_rows * THREADS_PER_ROW) { - const int row = tid / THREADS_PER_ROW; - const int thread_in_row = tid % THREADS_PER_ROW; - const unsigned mask = (valid_q_rows == BLOCK_M) - ? 0xFFFFFFFFU - : __activemask(); - const int row_leader = __ffs(mask) - 1; - - float* sS_row_f = sS + row * S_STRIDE + sub_start; - __half* sP_row_h = sP + row * p_stride; - - const int vec_cols = sub_valid_k_rows >> 2; - const int vecs_per_thread = - (vec_cols + THREADS_PER_ROW - 1) / THREADS_PER_ROW; - const int tail_start = vec_cols << 2; - - float thread_max = NEG_INF; - float4* sS_vec4 = reinterpret_cast(sS_row_f); - - #pragma unroll 4 - for (int j = 0; j < vecs_per_thread; ++j) { - int vc = thread_in_row + j * THREADS_PER_ROW; - if (vc < vec_cols) { - float4 v4 = sS_vec4[vc]; - thread_max = fmaxf( - thread_max, - fmaxf(fmaxf(v4.x, v4.y), fmaxf(v4.z, v4.w))); - } - } - - #pragma unroll 4 - for (int c = tail_start + thread_in_row; c < sub_valid_k_rows; - c += THREADS_PER_ROW) { - thread_max = fmaxf(thread_max, sS_row_f[c]); - } - - #pragma unroll - for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) { - thread_max = fmaxf( - thread_max, - __shfl_down_sync(mask, thread_max, o, THREADS_PER_ROW)); - } - - const float row_max = - __shfl_sync(mask, thread_max, row_leader, THREADS_PER_ROW); - const float old_max = sRowMax[row]; - const float new_max = fmaxf(old_max, row_max); - const float exp_diff = (D == 256) - ? expf(fmaxf(old_max - new_max, -80.0f)) - : exp2f(old_max - new_max); - - float thread_sum = 0.0f; - __half2 half_buffer[20]; - int vc_base = thread_in_row; - int h2_idx = 0; - int tail_col = -1; - __half tail_value = __float2half(0.f); - - #pragma unroll 4 - for (int j = 0; j < vecs_per_thread; ++j, - vc_base += THREADS_PER_ROW) { - if (vc_base < vec_cols) { - float4 v4 = sS_vec4[vc_base]; - - float e0 = (D == 256) - ? expf(fmaxf(v4.x - new_max, -80.0f)) - : exp2f(fmaxf(v4.x - new_max, EXP2_CLAMP)); - float e1 = (D == 256) - ? expf(fmaxf(v4.y - new_max, -80.0f)) - : exp2f(fmaxf(v4.y - new_max, EXP2_CLAMP)); - float e2 = (D == 256) - ? expf(fmaxf(v4.z - new_max, -80.0f)) - : exp2f(fmaxf(v4.z - new_max, EXP2_CLAMP)); - float e3 = (D == 256) - ? expf(fmaxf(v4.w - new_max, -80.0f)) - : exp2f(fmaxf(v4.w - new_max, EXP2_CLAMP)); - - thread_sum += (e0 + e1) + (e2 + e3); - - half_buffer[h2_idx++] = - __float22half2_rn(make_float2(e0, e1)); - half_buffer[h2_idx++] = - __float22half2_rn(make_float2(e2, e3)); - } - } - - #pragma unroll 4 - for (int c = tail_start + thread_in_row; c < sub_valid_k_rows; - c += THREADS_PER_ROW) { - float v = sS_row_f[c]; - float e = (D == 256) - ? expf(fmaxf(v - new_max, -80.0f)) - : exp2f(fmaxf(v - new_max, EXP2_CLAMP)); - thread_sum += e; - tail_col = c; - tail_value = __float2half_rn(e); - } - - #pragma unroll - for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) { - thread_sum += - __shfl_down_sync(mask, thread_sum, o, THREADS_PER_ROW); - } - - float row_sum = - __shfl_sync(mask, thread_sum, row_leader, THREADS_PER_ROW); - - if (thread_in_row == 0) { - sRowSum[row] = exp_diff * sRowSum[row] + row_sum; - sRowMax[row] = new_max; - } - - h2_idx = 0; - vc_base = thread_in_row; - __half2* sP_half2 = reinterpret_cast<__half2*>(sP_row_h); - - #pragma unroll 4 - for (int j = 0; j < vecs_per_thread; ++j, - vc_base += THREADS_PER_ROW) { - if (vc_base < vec_cols) { - int base_offset = vc_base * 2; - sP_half2[base_offset] = half_buffer[h2_idx++]; - sP_half2[base_offset + 1] = half_buffer[h2_idx++]; - } - } - - if (tail_col >= 0) { - sP_row_h[tail_col] = tail_value; - } - - #pragma unroll 4 - for (int c = tail_start + thread_in_row; c < p_tile_capacity; - c += THREADS_PER_ROW) { - if (c >= sub_valid_k_rows) { - sP_row_h[c] = __float2half(0.f); - } - } - - if (block_n > 0 || sub_start > 0) { - float* sO_row = sO + row * O_STRIDE; - float4* sO_vec = reinterpret_cast(sO_row); - const int o_vec_count = (O_STRIDE + 3) >> 2; - float scale = exp_diff; - - #pragma unroll 4 - for (int ov = thread_in_row; ov < o_vec_count; - ov += THREADS_PER_ROW) { - float4 v = sO_vec[ov]; - v.x *= scale; - v.y *= scale; - v.z *= scale; - v.w *= scale; - sO_vec[ov] = v; - } - } - } - __syncthreads(); - - if (scalar_pv) { - const int total_row_cols = valid_q_rows * D; - for (int idx = tid; idx < total_row_cols; - idx += THREADS_PER_BLOCK) { - const int row = idx / D; - const int d_col = idx - row * D; - const __half* sP_row_h = sP + row * p_stride; - float acc = 0.0f; - - #pragma unroll 1 - for (int k_idx = 0; k_idx < sub_valid_k_rows; ++k_idx) { - const float p_val = __half2float(sP_row_h[k_idx]); - const float v_val = __half2float( - sV[(sub_start + k_idx) * KV_STRIDE + d_col]); - acc += p_val * v_val; - } - sO[row * O_STRIDE + d_col] += acc; - } - __syncthreads(); - continue; - } - - const int num_tiles_k_pv = - (p_tile_capacity + WMMA_K - 1) / WMMA_K; - - for (int tile_idx = 0; tile_idx < tiles_per_warp_pv; ++tile_idx) { - const int global_tile_idx = warp_id * tiles_per_warp_pv + tile_idx; - if (global_tile_idx >= total_tiles_pv) break; - - const int tile_m_idx = global_tile_idx / num_tiles_n_pv; - const int tile_d_idx = global_tile_idx % num_tiles_n_pv; - - const int tile_m = tile_m_idx * WMMA_M; - const int tile_d = tile_d_idx * WMMA_N; - - if (tile_m >= valid_q_rows) continue; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - - load_matrix_sync( - acc_frag, - sO + tile_m * O_STRIDE + tile_d, - O_STRIDE, - mem_row_major); - - #pragma unroll - for (int tile_k = 0; tile_k < num_tiles_k_pv; ++tile_k) { - const int k_offset = tile_k * WMMA_K; - if (k_offset >= sub_valid_k_rows) break; - - load_matrix_sync( - a_frag, - sP + tile_m * p_stride + k_offset, - p_stride); - load_matrix_sync( - b_frag, - sV + (sub_start + k_offset) * KV_STRIDE + tile_d, - KV_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - store_matrix_sync( - sO + tile_m * O_STRIDE + tile_d, - acc_frag, - O_STRIDE, - mem_row_major); - } - __syncthreads(); +#pragma unroll + for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) { + thread_sum += __shfl_down_sync(mask, thread_sum, o, THREADS_PER_ROW); } - } - const int total_fp16_x4 = (valid_q_rows * D) / 4; - - for (int i = tid; i < total_fp16_x4; i += THREADS_PER_BLOCK) { - const int row = i / (D / 4); - const int col = (i % (D / 4)) * 4; - - const float sum_clamped = fmaxf(sRowSum[row], 1e-24f); - const float inv_sum = 1.0f / sum_clamped; - const float* sO_row = sO + row * O_STRIDE; - - const __half h0 = __float2half_rn(sO_row[col + 0] * inv_sum); - const __half h1 = __float2half_rn(sO_row[col + 1] * inv_sum); - const __half h2 = __float2half_rn(sO_row[col + 2] * inv_sum); - const __half h3 = __float2half_rn(sO_row[col + 3] * inv_sum); - - asm volatile( - "st.global.v4.u16 [%0], {%1, %2, %3, %4};" - : - : "l"(out_ptr + row * D + col), - "h"(__half_as_ushort(h0)), - "h"(__half_as_ushort(h1)), - "h"(__half_as_ushort(h2)), - "h"(__half_as_ushort(h3)) - : "memory" - ); - } + float row_sum = + __shfl_sync(mask, thread_sum, row_leader, THREADS_PER_ROW); - if (tid < valid_q_rows) { - const float sum = fmaxf(sRowSum[tid], 1e-24f); - softmax_lse_ptr[tid] = (D == 256) - ? sRowMax[tid] + logf(sum) - : (sRowMax[tid] + log2f(sum)) * LN2; - } -} + if (thread_in_row == 0) { + sRowSum[row] = exp_diff * sRowSum[row] + row_sum; + sRowMax[row] = new_max; + } -template -__global__ void __launch_bounds__(KernelConfig::THREADS_PER_BLOCK, 2) -flash_attention_qk_scores_kernel( - const __half* __restrict__ Q, - const __half* __restrict__ K, - float* __restrict__ Scores, - const int B, - const int H, - const int H_KV, - const int M, - const int N, - const float softmax_scale -) { - using Config = KernelConfig; - constexpr int BLOCK_M = Config::BLOCK_M; - constexpr int BLOCK_N = Config::BLOCK_N; - constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; - constexpr int WARPS_PER_BLOCK = Config::WARPS_PER_BLOCK; - constexpr int Q_STRIDE = Config::Q_STRIDE; - constexpr int KV_STRIDE = Config::KV_STRIDE; - constexpr int S_STRIDE = Config::S_STRIDE; - constexpr int PER_UINT4 = Config::PER_UINT4; - const float NEG_INF = -1e30f; - - const int batch_head_id = blockIdx.z; - if (batch_head_id >= B * H) return; - const int batch_id = batch_head_id / H; - const int q_head_id = batch_head_id % H; - const int kv_group_size = H / H_KV; - const int kv_head_id = q_head_id / kv_group_size; - - const int block_m = blockIdx.x; - const int block_n = blockIdx.y; - const int start_row = block_m * BLOCK_M; - const int start_col = block_n * BLOCK_N; - if (start_row >= M || start_col >= N) return; + h2_idx = 0; + vc_base = thread_in_row; + __half2* sP_half2 = reinterpret_cast<__half2*>(sP_row_h); + +#pragma unroll 4 + for (int j = 0; j < vecs_per_thread; ++j, vc_base += THREADS_PER_ROW) { + if (vc_base < vec_cols) { + int base_offset = vc_base * 2; + sP_half2[base_offset] = half_buffer[h2_idx++]; + sP_half2[base_offset + 1] = half_buffer[h2_idx++]; + } + } - const int valid_q_rows = min(BLOCK_M, M - start_row); - const int valid_k_rows = min(BLOCK_N, N - start_col); - const int causal_q_offset = max(N - M, 0); + if (tail_col >= 0) { + sP_row_h[tail_col] = tail_value; + } - if constexpr (IS_CAUSAL) { - if (start_col >= causal_q_offset + start_row + valid_q_rows) { - return; +#pragma unroll 4 + for (int c = tail_start + thread_in_row; c < p_tile_capacity; + c += THREADS_PER_ROW) { + if (c >= sub_valid_k_rows) { + sP_row_h[c] = __float2half(0.f); + } } - } - const int tid = threadIdx.x; - const int warp_id = tid / MAX_THREADS_PER_WARP; - const int lane_id = tid % MAX_THREADS_PER_WARP; + if (block_n > 0 || sub_start > 0) { + float* sO_row = sO + row * O_STRIDE; + float4* sO_vec = reinterpret_cast(sO_row); + const int o_vec_count = (O_STRIDE + 3) >> 2; + float scale = exp_diff; + +#pragma unroll 4 + for (int ov = thread_in_row; ov < o_vec_count; + ov += THREADS_PER_ROW) { + float4 v = sO_vec[ov]; + v.x *= scale; + v.y *= scale; + v.z *= scale; + v.w *= scale; + sO_vec[ov] = v; + } + } + } + __syncthreads(); + + if (scalar_pv) { + const int total_row_cols = valid_q_rows * D; + for (int idx = tid; idx < total_row_cols; idx += THREADS_PER_BLOCK) { + const int row = idx / D; + const int d_col = idx - row * D; + const __half* sP_row_h = sP + row * p_stride; + float acc = 0.0f; + +#pragma unroll 1 + for (int k_idx = 0; k_idx < sub_valid_k_rows; ++k_idx) { + const float p_val = __half2float(sP_row_h[k_idx]); + const float v_val = + __half2float(sV[(sub_start + k_idx) * KV_STRIDE + d_col]); + acc += p_val * v_val; + } + sO[row * O_STRIDE + d_col] += acc; + } + __syncthreads(); + continue; + } - const size_t q_head_linear = (size_t)batch_id * H + q_head_id; - const size_t kv_head_linear = (size_t)batch_id * H_KV + kv_head_id; - const __half* q_ptr = Q + q_head_linear * M * D + start_row * D; - const __half* k_ptr = K + kv_head_linear * N * D + start_col * D; - float* score_ptr = Scores + q_head_linear * M * N + start_row * N + start_col; + const int num_tiles_k_pv = (p_tile_capacity + WMMA_K - 1) / WMMA_K; - extern __shared__ char smem_raw[]; - init_smem(smem_raw); - auto& smem = *reinterpret_cast(smem_raw); + for (int tile_idx = 0; tile_idx < tiles_per_warp_pv; ++tile_idx) { + const int global_tile_idx = warp_id * tiles_per_warp_pv + tile_idx; + if (global_tile_idx >= total_tiles_pv) break; - __half* sQ = smem.q; - __half* sK = smem.reuse_kv.k; - float* sS = smem.reuse_sp.s; + const int tile_m_idx = global_tile_idx / num_tiles_n_pv; + const int tile_d_idx = global_tile_idx % num_tiles_n_pv; - const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; - const int q_stride_uint4 = (Q_STRIDE + PER_UINT4 - 1) / PER_UINT4; - const int kv_stride_uint4 = (KV_STRIDE + PER_UINT4 - 1) / PER_UINT4; + const int tile_m = tile_m_idx * WMMA_M; + const int tile_d = tile_d_idx * WMMA_N; - const uint4* q_vec = reinterpret_cast(q_ptr); - uint4* sQ_vec = reinterpret_cast(sQ); + if (tile_m >= valid_q_rows) continue; - #pragma unroll 4 - for (int idx = tid; idx < (valid_q_rows * d_stride_uint4); - idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - uint4 q_val = make_uint4(0, 0, 0, 0); - if (row < valid_q_rows && vec_col < d_stride_uint4) { - q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); - } - sQ_vec[row * q_stride_uint4 + vec_col] = q_val; - } + fragment a_frag; + fragment b_frag; + fragment acc_frag; - const uint4* k_vec = reinterpret_cast(k_ptr); - uint4* sK_vec = reinterpret_cast(sK); + load_matrix_sync(acc_frag, sO + tile_m * O_STRIDE + tile_d, O_STRIDE, + mem_row_major); - #pragma unroll 2 - for (int idx = tid; idx < (valid_k_rows * d_stride_uint4); - idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - uint4 k_val = make_uint4(0, 0, 0, 0); - if (row < valid_k_rows && vec_col < d_stride_uint4) { - k_val = __ldg(&k_vec[row * d_stride_uint4 + vec_col]); - } - sK_vec[row * kv_stride_uint4 + vec_col] = k_val; - } - __syncthreads(); +#pragma unroll + for (int tile_k = 0; tile_k < num_tiles_k_pv; ++tile_k) { + const int k_offset = tile_k * WMMA_K; + if (k_offset >= sub_valid_k_rows) break; - if constexpr (D == 256) { - compute_qk_scores_scalar_fp32( - sQ, sK, sS, valid_q_rows, valid_k_rows, start_row, start_col, - causal_q_offset, softmax_scale, -1, -1); - } else { - const int num_tiles_m_qk = (BLOCK_M + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_qk = (BLOCK_N + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_qk = (D + WMMA_K - 1) / WMMA_K; - const int total_tiles_qk = num_tiles_m_qk * num_tiles_n_qk; - const int tiles_per_warp_qk = - (total_tiles_qk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - const unsigned row_causal = - (lane_id & 0b1) + ((lane_id >> 2) & 0b1) * 8 - + ((lane_id >> 4) & 0b1) * 4; - const unsigned col_causal = - ((lane_id >> 1) & 0b1) * 2 + ((lane_id >> 3) & 0b1) * 8; - - for (int tile_idx = 0; tile_idx < tiles_per_warp_qk; ++tile_idx) { - const int global_tile_idx = warp_id * tiles_per_warp_qk + tile_idx; - if (global_tile_idx >= total_tiles_qk) break; - - const int tile_m_idx = global_tile_idx / num_tiles_n_qk; - const int tile_n_idx = global_tile_idx % num_tiles_n_qk; - - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; - - if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - fill_fragment(acc_frag, 0.0f); - - #pragma unroll - for (int k_tile = 0; k_tile < num_tiles_k_qk; ++k_tile) { - const int k_offset = k_tile * WMMA_K; - if (k_offset >= D) break; - - load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, - Q_STRIDE); - load_matrix_sync(b_frag, sK + tile_n * KV_STRIDE + k_offset, - KV_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - - if constexpr (IS_CAUSAL) { - #pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - const unsigned col = - col_causal + (i & 0b1) + ((i >> 2) & 0b1) * 4; - const unsigned row = row_causal + ((i >> 1) & 0b1) * 2; - - const int global_m = start_row + tile_m + row; - const int global_n = start_col + tile_n + col; - const int global_q_pos = global_m + causal_q_offset; - - const bool is_valid = - (global_m < start_row + valid_q_rows) && - (global_n < start_col + valid_k_rows); - - acc_frag.x[i] = is_valid - ? ((global_n > global_q_pos) - ? NEG_INF - : acc_frag.x[i] * softmax_scale) - : NEG_INF; - } - } else { - #pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - acc_frag.x[i] *= softmax_scale; - } - } - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, - S_STRIDE, mem_row_major); + load_matrix_sync(a_frag, sP + tile_m * p_stride + k_offset, p_stride); + load_matrix_sync(b_frag, + sV + (sub_start + k_offset) * KV_STRIDE + tile_d, + KV_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); } + store_matrix_sync(sO + tile_m * O_STRIDE + tile_d, acc_frag, O_STRIDE, + mem_row_major); + } + __syncthreads(); } - __syncthreads(); + } + + const int total_fp16_x4 = (valid_q_rows * D) / 4; + + for (int i = tid; i < total_fp16_x4; i += THREADS_PER_BLOCK) { + const int row = i / (D / 4); + const int col = (i % (D / 4)) * 4; + + const float sum_clamped = fmaxf(sRowSum[row], 1e-24f); + const float inv_sum = 1.0f / sum_clamped; + const float* sO_row = sO + row * O_STRIDE; + + const __half h0 = __float2half_rn(sO_row[col + 0] * inv_sum); + const __half h1 = __float2half_rn(sO_row[col + 1] * inv_sum); + const __half h2 = __float2half_rn(sO_row[col + 2] * inv_sum); + const __half h3 = __float2half_rn(sO_row[col + 3] * inv_sum); + + asm volatile("st.global.v4.u16 [%0], {%1, %2, %3, %4};" + : + : "l"(out_ptr + row * D + col), "h"(__half_as_ushort(h0)), + "h"(__half_as_ushort(h1)), "h"(__half_as_ushort(h2)), + "h"(__half_as_ushort(h3)) + : "memory"); + } + + if (tid < valid_q_rows) { + const float sum = fmaxf(sRowSum[tid], 1e-24f); + softmax_lse_ptr[tid] = (D == 256) ? sRowMax[tid] + logf(sum) + : (sRowMax[tid] + log2f(sum)) * LN2; + } +} - for (int idx = tid; idx < valid_q_rows * valid_k_rows; - idx += THREADS_PER_BLOCK) { - const int row = idx / valid_k_rows; - const int col = idx % valid_k_rows; - score_ptr[row * N + col] = sS[row * S_STRIDE + col]; +template +__global__ void __launch_bounds__(KernelConfig::THREADS_PER_BLOCK, 2) + flash_attention_qk_scores_kernel(const __half* __restrict__ Q, + const __half* __restrict__ K, + float* __restrict__ Scores, const int B, + const int H, const int H_KV, const int M, + const int N, const float softmax_scale) { + using Config = KernelConfig; + constexpr int BLOCK_M = Config::BLOCK_M; + constexpr int BLOCK_N = Config::BLOCK_N; + constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; + constexpr int WARPS_PER_BLOCK = Config::WARPS_PER_BLOCK; + constexpr int Q_STRIDE = Config::Q_STRIDE; + constexpr int KV_STRIDE = Config::KV_STRIDE; + constexpr int S_STRIDE = Config::S_STRIDE; + constexpr int PER_UINT4 = Config::PER_UINT4; + const float NEG_INF = -1e30f; + + const int batch_head_id = blockIdx.z; + if (batch_head_id >= B * H) return; + const int batch_id = batch_head_id / H; + const int q_head_id = batch_head_id % H; + const int kv_group_size = H / H_KV; + const int kv_head_id = q_head_id / kv_group_size; + + const int block_m = blockIdx.x; + const int block_n = blockIdx.y; + const int start_row = block_m * BLOCK_M; + const int start_col = block_n * BLOCK_N; + if (start_row >= M || start_col >= N) return; + + const int valid_q_rows = min(BLOCK_M, M - start_row); + const int valid_k_rows = min(BLOCK_N, N - start_col); + const int causal_q_offset = max(N - M, 0); + + if constexpr (IS_CAUSAL) { + if (start_col >= causal_q_offset + start_row + valid_q_rows) { + return; + } + } + + const int tid = threadIdx.x; + const int warp_id = tid / MAX_THREADS_PER_WARP; + const int lane_id = tid % MAX_THREADS_PER_WARP; + + const size_t q_head_linear = (size_t)batch_id * H + q_head_id; + const size_t kv_head_linear = (size_t)batch_id * H_KV + kv_head_id; + const __half* q_ptr = Q + q_head_linear * M * D + start_row * D; + const __half* k_ptr = K + kv_head_linear * N * D + start_col * D; + float* score_ptr = Scores + q_head_linear * M * N + start_row * N + start_col; + + extern __shared__ char smem_raw[]; + init_smem(smem_raw); + auto& smem = *reinterpret_cast(smem_raw); + + __half* sQ = smem.q; + __half* sK = smem.reuse_kv.k; + float* sS = smem.reuse_sp.s; + + const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; + const int q_stride_uint4 = (Q_STRIDE + PER_UINT4 - 1) / PER_UINT4; + const int kv_stride_uint4 = (KV_STRIDE + PER_UINT4 - 1) / PER_UINT4; + + const uint4* q_vec = reinterpret_cast(q_ptr); + uint4* sQ_vec = reinterpret_cast(sQ); + +#pragma unroll 4 + for (int idx = tid; idx < (valid_q_rows * d_stride_uint4); + idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + uint4 q_val = make_uint4(0, 0, 0, 0); + if (row < valid_q_rows && vec_col < d_stride_uint4) { + q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); + } + sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + } + + const uint4* k_vec = reinterpret_cast(k_ptr); + uint4* sK_vec = reinterpret_cast(sK); + +#pragma unroll 2 + for (int idx = tid; idx < (valid_k_rows * d_stride_uint4); + idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + uint4 k_val = make_uint4(0, 0, 0, 0); + if (row < valid_k_rows && vec_col < d_stride_uint4) { + k_val = __ldg(&k_vec[row * d_stride_uint4 + vec_col]); + } + sK_vec[row * kv_stride_uint4 + vec_col] = k_val; + } + __syncthreads(); + + if constexpr (D == 256) { + compute_qk_scores_scalar_fp32( + sQ, sK, sS, valid_q_rows, valid_k_rows, start_row, start_col, + causal_q_offset, softmax_scale, -1, -1); + } else { + const int num_tiles_m_qk = (BLOCK_M + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_qk = (BLOCK_N + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_qk = (D + WMMA_K - 1) / WMMA_K; + const int total_tiles_qk = num_tiles_m_qk * num_tiles_n_qk; + const int tiles_per_warp_qk = + (total_tiles_qk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + const unsigned row_causal = (lane_id & 0b1) + ((lane_id >> 2) & 0b1) * 8 + + ((lane_id >> 4) & 0b1) * 4; + const unsigned col_causal = + ((lane_id >> 1) & 0b1) * 2 + ((lane_id >> 3) & 0b1) * 8; + + for (int tile_idx = 0; tile_idx < tiles_per_warp_qk; ++tile_idx) { + const int global_tile_idx = warp_id * tiles_per_warp_qk + tile_idx; + if (global_tile_idx >= total_tiles_qk) break; + + const int tile_m_idx = global_tile_idx / num_tiles_n_qk; + const int tile_n_idx = global_tile_idx % num_tiles_n_qk; + + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; + + if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + fill_fragment(acc_frag, 0.0f); + +#pragma unroll + for (int k_tile = 0; k_tile < num_tiles_k_qk; ++k_tile) { + const int k_offset = k_tile * WMMA_K; + if (k_offset >= D) break; + + load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, Q_STRIDE); + load_matrix_sync(b_frag, sK + tile_n * KV_STRIDE + k_offset, KV_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + if constexpr (IS_CAUSAL) { +#pragma unroll + for (int i = 0; i < acc_frag.num_elements; ++i) { + const unsigned col = col_causal + (i & 0b1) + ((i >> 2) & 0b1) * 4; + const unsigned row = row_causal + ((i >> 1) & 0b1) * 2; + + const int global_m = start_row + tile_m + row; + const int global_n = start_col + tile_n + col; + const int global_q_pos = global_m + causal_q_offset; + + const bool is_valid = (global_m < start_row + valid_q_rows) && + (global_n < start_col + valid_k_rows); + + acc_frag.x[i] = is_valid ? ((global_n > global_q_pos) + ? NEG_INF + : acc_frag.x[i] * softmax_scale) + : NEG_INF; + } + } else { +#pragma unroll + for (int i = 0; i < acc_frag.num_elements; ++i) { + acc_frag.x[i] *= softmax_scale; + } + } + store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, + mem_row_major); } + } + __syncthreads(); + + for (int idx = tid; idx < valid_q_rows * valid_k_rows; + idx += THREADS_PER_BLOCK) { + const int row = idx / valid_k_rows; + const int col = idx % valid_k_rows; + score_ptr[row * N + col] = sS[row * S_STRIDE + col]; + } } -template +template void launcher_flash_attention_forward_impl( - const torch::Tensor& Q, - const torch::Tensor& K, - const torch::Tensor& V, - torch::Tensor& Out, - torch::Tensor& softmax_lse, - float softmax_scale, - bool is_causal, - bool scalar_pv, - int window_size_left, - int window_size_right, - cudaStream_t stream -) { - using Config = KernelConfig; - - const int B = Q.size(0); - const int H = Q.size(1); - const int H_KV = K.size(1); - const int M = Q.size(2); - const int N = K.size(2); - - const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; - const dim3 grid(grid_x, 1, B * H); - const dim3 block(Config::THREADS_PER_BLOCK); - const size_t smem = Config::TOTAL_SMEM; - - TORCH_CHECK(smem <= MAX_SMEM_PER_SM, "Shared memory exceeds 96KB: ", smem, " bytes"); - - auto kernel = is_causal ? - (void*)flash_attention_forward_kernel : - (void*)flash_attention_forward_kernel; - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - - if (is_causal) { - flash_attention_forward_kernel + const torch::Tensor& Q, const torch::Tensor& K, const torch::Tensor& V, + torch::Tensor& Out, torch::Tensor& softmax_lse, float softmax_scale, + bool is_causal, bool scalar_pv, int window_size_left, int window_size_right, + cudaStream_t stream) { + using Config = KernelConfig; + + const int B = Q.size(0); + const int H = Q.size(1); + const int H_KV = K.size(1); + const int M = Q.size(2); + const int N = K.size(2); + + const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; + const dim3 grid(grid_x, 1, B * H); + const dim3 block(Config::THREADS_PER_BLOCK); + const size_t smem = Config::TOTAL_SMEM; + + TORCH_CHECK(smem <= MAX_SMEM_PER_SM, "Shared memory exceeds 96KB: ", smem, + " bytes"); + + auto kernel = + is_causal + ? (void*) + flash_attention_forward_kernel + : (void*)flash_attention_forward_kernel; + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + if (is_causal) { + flash_attention_forward_kernel <<>>( reinterpret_cast(Q.data_ptr()), reinterpret_cast(K.data_ptr()), reinterpret_cast(V.data_ptr()), reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), - B, H, H_KV, M, N, softmax_scale, scalar_pv, - window_size_left, window_size_right - ); - } else { - flash_attention_forward_kernel + softmax_lse.data_ptr(), B, H, H_KV, M, N, softmax_scale, + scalar_pv, window_size_left, window_size_right); + } else { + flash_attention_forward_kernel <<>>( reinterpret_cast(Q.data_ptr()), reinterpret_cast(K.data_ptr()), reinterpret_cast(V.data_ptr()), reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), - B, H, H_KV, M, N, softmax_scale, scalar_pv, - window_size_left, window_size_right - ); - } + softmax_lse.data_ptr(), B, H, H_KV, M, N, softmax_scale, + scalar_pv, window_size_left, window_size_right); + } } -template +template void launcher_flash_attention_forward( - const torch::Tensor& Q, - const torch::Tensor& K, - const torch::Tensor& V, - torch::Tensor& Out, - torch::Tensor& softmax_lse, - float softmax_scale, - bool is_causal, - bool scalar_pv, - int window_size_left, - int window_size_right, - cudaStream_t stream -) { - if constexpr (D == 256) { - const bool use_wmma_qk = - env_flag_enabled_default_on("VLLM_FLASH_V100_DENSE_D256_WMMA_QK"); - if (env_flag_enabled("VLLM_FLASH_V100_DENSE_D256_LOW_SMEM")) { - if (use_wmma_qk) { - launcher_flash_attention_forward_impl( - Q, K, V, Out, softmax_lse, softmax_scale, is_causal, - scalar_pv, window_size_left, window_size_right, stream); - } else { - launcher_flash_attention_forward_impl( - Q, K, V, Out, softmax_lse, softmax_scale, is_causal, - scalar_pv, window_size_left, window_size_right, stream); - } - return; - } - if (use_wmma_qk) { - launcher_flash_attention_forward_impl( - Q, K, V, Out, softmax_lse, softmax_scale, is_causal, - scalar_pv, window_size_left, window_size_right, stream); - return; - } + const torch::Tensor& Q, const torch::Tensor& K, const torch::Tensor& V, + torch::Tensor& Out, torch::Tensor& softmax_lse, float softmax_scale, + bool is_causal, bool scalar_pv, int window_size_left, int window_size_right, + cudaStream_t stream) { + if constexpr (D == 256) { + const bool use_wmma_qk = + env_flag_enabled_default_on("VLLM_FLASH_V100_DENSE_D256_WMMA_QK"); + if (env_flag_enabled("VLLM_FLASH_V100_DENSE_D256_LOW_SMEM")) { + if (use_wmma_qk) { + launcher_flash_attention_forward_impl( + Q, K, V, Out, softmax_lse, softmax_scale, is_causal, scalar_pv, + window_size_left, window_size_right, stream); + } else { + launcher_flash_attention_forward_impl( + Q, K, V, Out, softmax_lse, softmax_scale, is_causal, scalar_pv, + window_size_left, window_size_right, stream); + } + return; } - launcher_flash_attention_forward_impl( - Q, K, V, Out, softmax_lse, softmax_scale, is_causal, scalar_pv, - window_size_left, window_size_right, stream); -} - -template -void launcher_flash_attention_qk_scores( - const torch::Tensor& Q, - const torch::Tensor& K, - torch::Tensor& Scores, - float softmax_scale, - bool is_causal, - cudaStream_t stream -) { - using Config = KernelConfig; - - const int B = Q.size(0); - const int H = Q.size(1); - const int M = Q.size(2); - const int H_KV = K.size(1); - const int N = K.size(2); - - const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; - const int grid_y = (N + Config::BLOCK_N - 1) / Config::BLOCK_N; - const dim3 grid(grid_x, grid_y, B * H); - const dim3 block(Config::THREADS_PER_BLOCK); - const size_t smem = Config::TOTAL_SMEM; - - TORCH_CHECK(smem <= MAX_SMEM_PER_SM, "Shared memory exceeds 96KB: ", smem, " bytes"); - - auto kernel = is_causal ? - (void*)flash_attention_qk_scores_kernel : - (void*)flash_attention_qk_scores_kernel; - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - - if (is_causal) { - flash_attention_qk_scores_kernel<<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K.data_ptr()), - Scores.data_ptr(), - B, H, H_KV, M, N, softmax_scale - ); - } else { - flash_attention_qk_scores_kernel<<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K.data_ptr()), - Scores.data_ptr(), - B, H, H_KV, M, N, softmax_scale - ); + if (use_wmma_qk) { + launcher_flash_attention_forward_impl( + Q, K, V, Out, softmax_lse, softmax_scale, is_causal, scalar_pv, + window_size_left, window_size_right, stream); + return; } + } + launcher_flash_attention_forward_impl( + Q, K, V, Out, softmax_lse, softmax_scale, is_causal, scalar_pv, + window_size_left, window_size_right, stream); } -at::Tensor flash_attention_qk_scores( - const at::Tensor& q, - const at::Tensor& k, - const float softmax_scale, - const bool is_causal -) { - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - TORCH_CHECK(k.dtype() == torch::kFloat16, "k must be fp16"); - TORCH_CHECK(q.is_cuda() && k.is_cuda(), "Tensors must be on CUDA"); - TORCH_CHECK(q.stride(-1) == 1 && k.stride(-1) == 1, - "Last dim must be contiguous"); - TORCH_CHECK(q.dim() == 4 && k.dim() == 4, - "q/k must have shape [B, H, T, D]"); - - const int B = q.size(0); - const int H = q.size(1); - const int M = q.size(2); - const int D = q.size(3); - const int H_KV = k.size(1); - const int N = k.size(2); - - TORCH_CHECK(k.size(0) == B, "K batch size must match Q"); - TORCH_CHECK(k.size(3) == D, "K head_dim must match Q"); - TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, - "D must be even, <=256, multiple of 8"); - TORCH_CHECK(H_KV > 0, "num_kv_heads must be positive"); - TORCH_CHECK(H % H_KV == 0, - "num_attention_heads must be divisible by num_kv_heads"); - - auto props = at::cuda::getCurrentDeviceProperties(); - bool sm70 = props->major == 7 && props->minor == 0; - TORCH_CHECK(sm70, "Kernel supports only Volta GPUs."); - - auto scores = torch::full({B, H, M, N}, -1e30f, - torch::dtype(torch::kFloat32).device(q.device())); - auto stream = at::cuda::getCurrentCUDAStream().stream(); - - switch (D) { - case 16: launcher_flash_attention_qk_scores<16>(q, k, scores, softmax_scale, is_causal, stream); break; - case 32: launcher_flash_attention_qk_scores<32>(q, k, scores, softmax_scale, is_causal, stream); break; - case 64: launcher_flash_attention_qk_scores<64>(q, k, scores, softmax_scale, is_causal, stream); break; - case 128: launcher_flash_attention_qk_scores<128>(q, k, scores, softmax_scale, is_causal, stream); break; - case 256: launcher_flash_attention_qk_scores<256>(q, k, scores, softmax_scale, is_causal, stream); break; - default: TORCH_CHECK(false, "Unsupported D: ", D); - } +template +void launcher_flash_attention_qk_scores(const torch::Tensor& Q, + const torch::Tensor& K, + torch::Tensor& Scores, + float softmax_scale, bool is_causal, + cudaStream_t stream) { + using Config = KernelConfig; + + const int B = Q.size(0); + const int H = Q.size(1); + const int M = Q.size(2); + const int H_KV = K.size(1); + const int N = K.size(2); + + const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; + const int grid_y = (N + Config::BLOCK_N - 1) / Config::BLOCK_N; + const dim3 grid(grid_x, grid_y, B * H); + const dim3 block(Config::THREADS_PER_BLOCK); + const size_t smem = Config::TOTAL_SMEM; + + TORCH_CHECK(smem <= MAX_SMEM_PER_SM, "Shared memory exceeds 96KB: ", smem, + " bytes"); + + auto kernel = is_causal ? (void*)flash_attention_qk_scores_kernel + : (void*)flash_attention_qk_scores_kernel; + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + if (is_causal) { + flash_attention_qk_scores_kernel<<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K.data_ptr()), Scores.data_ptr(), + B, H, H_KV, M, N, softmax_scale); + } else { + flash_attention_qk_scores_kernel<<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K.data_ptr()), Scores.data_ptr(), + B, H, H_KV, M, N, softmax_scale); + } +} - return scores; +at::Tensor flash_attention_qk_scores(const at::Tensor& q, const at::Tensor& k, + const float softmax_scale, + const bool is_causal) { + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + TORCH_CHECK(k.dtype() == torch::kFloat16, "k must be fp16"); + TORCH_CHECK(q.is_cuda() && k.is_cuda(), "Tensors must be on CUDA"); + TORCH_CHECK(q.stride(-1) == 1 && k.stride(-1) == 1, + "Last dim must be contiguous"); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "q/k must have shape [B, H, T, D]"); + + const int B = q.size(0); + const int H = q.size(1); + const int M = q.size(2); + const int D = q.size(3); + const int H_KV = k.size(1); + const int N = k.size(2); + + TORCH_CHECK(k.size(0) == B, "K batch size must match Q"); + TORCH_CHECK(k.size(3) == D, "K head_dim must match Q"); + TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, + "D must be even, <=256, multiple of 8"); + TORCH_CHECK(H_KV > 0, "num_kv_heads must be positive"); + TORCH_CHECK(H % H_KV == 0, + "num_attention_heads must be divisible by num_kv_heads"); + + auto props = at::cuda::getCurrentDeviceProperties(); + bool sm70 = props->major == 7 && props->minor == 0; + TORCH_CHECK(sm70, "Kernel supports only Volta GPUs."); + + auto scores = torch::full({B, H, M, N}, -1e30f, + torch::dtype(torch::kFloat32).device(q.device())); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + switch (D) { + case 16: + launcher_flash_attention_qk_scores<16>(q, k, scores, softmax_scale, + is_causal, stream); + break; + case 32: + launcher_flash_attention_qk_scores<32>(q, k, scores, softmax_scale, + is_causal, stream); + break; + case 64: + launcher_flash_attention_qk_scores<64>(q, k, scores, softmax_scale, + is_causal, stream); + break; + case 128: + launcher_flash_attention_qk_scores<128>(q, k, scores, softmax_scale, + is_causal, stream); + break; + case 256: + launcher_flash_attention_qk_scores<256>(q, k, scores, softmax_scale, + is_causal, stream); + break; + default: + TORCH_CHECK(false, "Unsupported D: ", D); + } + + return scores; } std::vector flash_attention_forward( - at::Tensor& q, - const at::Tensor& k, - const at::Tensor& v, - std::optional& out_, - std::optional& alibi_slopes_, - const float p_dropout, - const float softmax_scale, - bool is_causal, - int window_size_left, - int window_size_right, - const float softcap, - const bool return_softmax, - std::optional gen_ -) { - - TORCH_CHECK(!alibi_slopes_.has_value(), "alibi_slopes not supported"); - TORCH_CHECK(p_dropout == 0.f, "dropout not supported"); - TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, - "window sizes must be >= -1"); - TORCH_CHECK(softcap == 0.f, "softcap not supported"); - TORCH_CHECK(!return_softmax, "return_softmax not supported"); - TORCH_CHECK(!gen_.has_value(), "Generator not supported"); - - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - TORCH_CHECK(k.dtype() == torch::kFloat16, "k must be fp16"); - TORCH_CHECK(v.dtype() == torch::kFloat16, "v must be fp16"); - TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "Tensors must be on CUDA"); - TORCH_CHECK(q.stride(-1) == 1 && k.stride(-1) == 1 && v.stride(-1) == 1, "Last dim must be contiguous"); - - const auto sizes = q.sizes(); - const int B = sizes[0], H = sizes[1], M = sizes[2], D = sizes[3]; - const int H_KV = k.size(1); - const int N = k.size(2); - TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, "D must be even, <=256, multiple of 8"); - TORCH_CHECK(H_KV > 0, "num_kv_heads must be positive"); - TORCH_CHECK(H % H_KV == 0, "num_attention_heads must be divisible by num_kv_heads"); - TORCH_CHECK(k.size(0) == B && v.size(0) == B, "K/V batch size must match Q"); - TORCH_CHECK(v.size(1) == H_KV, "K/V num_heads must match"); - TORCH_CHECK(v.size(2) == N, "K/V sequence length must match"); - TORCH_CHECK(k.size(3) == D && v.size(3) == D, "K/V head_dim must match Q"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); - TORCH_CHECK(out_fp16.dtype() == torch::kFloat16, "out must be fp16"); - auto softmax_lse = torch::zeros({B, H, M}, torch::dtype(torch::kFloat32).device(q.device())); - TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, "softmax_lse must be fp32"); - - auto stream = at::cuda::getCurrentCUDAStream().stream(); - auto props = at::cuda::getCurrentDeviceProperties(); - bool sm70 = props->major == 7 && props->minor == 0; - TORCH_CHECK(sm70, "Kernel supports only Volta GPUs."); - const char* scalar_pv_env = std::getenv("VLLM_FLASH_V100_PREFILL_SCALAR_PV"); - const bool scalar_pv = scalar_pv_env != nullptr && scalar_pv_env[0] != '\0' && - scalar_pv_env[0] != '0'; - - switch (D) { - case 16: launcher_flash_attention_forward<16>(q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, scalar_pv, window_size_left, window_size_right, stream); break; - case 32: launcher_flash_attention_forward<32>(q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, scalar_pv, window_size_left, window_size_right, stream); break; - case 64: launcher_flash_attention_forward<64>(q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, scalar_pv, window_size_left, window_size_right, stream); break; - case 128: launcher_flash_attention_forward<128>(q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, scalar_pv, window_size_left, window_size_right, stream); break; - case 256: launcher_flash_attention_forward<256>(q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, scalar_pv, window_size_left, window_size_right, stream); break; - default: TORCH_CHECK(false, "Unsupported D: ", D); - } - - auto p = torch::zeros({0}, q.options()); - auto rng_state = torch::zeros({2}, torch::dtype(torch::kInt64).device(q.device())); - return {out_fp16, softmax_lse, p, rng_state}; + at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out_, std::optional& alibi_slopes_, + const float p_dropout, const float softmax_scale, bool is_causal, + int window_size_left, int window_size_right, const float softcap, + const bool return_softmax, std::optional gen_) { + TORCH_CHECK(!alibi_slopes_.has_value(), "alibi_slopes not supported"); + TORCH_CHECK(p_dropout == 0.f, "dropout not supported"); + TORCH_CHECK(window_size_left >= -1 && window_size_right >= -1, + "window sizes must be >= -1"); + TORCH_CHECK(softcap == 0.f, "softcap not supported"); + TORCH_CHECK(!return_softmax, "return_softmax not supported"); + TORCH_CHECK(!gen_.has_value(), "Generator not supported"); + + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + TORCH_CHECK(k.dtype() == torch::kFloat16, "k must be fp16"); + TORCH_CHECK(v.dtype() == torch::kFloat16, "v must be fp16"); + TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), + "Tensors must be on CUDA"); + TORCH_CHECK(q.stride(-1) == 1 && k.stride(-1) == 1 && v.stride(-1) == 1, + "Last dim must be contiguous"); + + const auto sizes = q.sizes(); + const int B = sizes[0], H = sizes[1], M = sizes[2], D = sizes[3]; + const int H_KV = k.size(1); + const int N = k.size(2); + TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, + "D must be even, <=256, multiple of 8"); + TORCH_CHECK(H_KV > 0, "num_kv_heads must be positive"); + TORCH_CHECK(H % H_KV == 0, + "num_attention_heads must be divisible by num_kv_heads"); + TORCH_CHECK(k.size(0) == B && v.size(0) == B, "K/V batch size must match Q"); + TORCH_CHECK(v.size(1) == H_KV, "K/V num_heads must match"); + TORCH_CHECK(v.size(2) == N, "K/V sequence length must match"); + TORCH_CHECK(k.size(3) == D && v.size(3) == D, "K/V head_dim must match Q"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); + TORCH_CHECK(out_fp16.dtype() == torch::kFloat16, "out must be fp16"); + auto softmax_lse = + torch::zeros({B, H, M}, torch::dtype(torch::kFloat32).device(q.device())); + TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, + "softmax_lse must be fp32"); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + auto props = at::cuda::getCurrentDeviceProperties(); + bool sm70 = props->major == 7 && props->minor == 0; + TORCH_CHECK(sm70, "Kernel supports only Volta GPUs."); + const char* scalar_pv_env = std::getenv("VLLM_FLASH_V100_PREFILL_SCALAR_PV"); + const bool scalar_pv = scalar_pv_env != nullptr && scalar_pv_env[0] != '\0' && + scalar_pv_env[0] != '0'; + + switch (D) { + case 16: + launcher_flash_attention_forward<16>( + q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, scalar_pv, + window_size_left, window_size_right, stream); + break; + case 32: + launcher_flash_attention_forward<32>( + q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, scalar_pv, + window_size_left, window_size_right, stream); + break; + case 64: + launcher_flash_attention_forward<64>( + q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, scalar_pv, + window_size_left, window_size_right, stream); + break; + case 128: + launcher_flash_attention_forward<128>( + q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, scalar_pv, + window_size_left, window_size_right, stream); + break; + case 256: + launcher_flash_attention_forward<256>( + q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, scalar_pv, + window_size_left, window_size_right, stream); + break; + default: + TORCH_CHECK(false, "Unsupported D: ", D); + } + + auto p = torch::zeros({0}, q.options()); + auto rng_state = + torch::zeros({2}, torch::dtype(torch::kInt64).device(q.device())); + return {out_fp16, softmax_lse, p, rng_state}; } diff --git a/flash-attention-v100/kernel/fused_mha_forward_paged.cu b/flash-attention-v100/kernel/fused_mha_forward_paged.cu index 538aaad265..0a5a55f406 100644 --- a/flash-attention-v100/kernel/fused_mha_forward_paged.cu +++ b/flash-attention-v100/kernel/fused_mha_forward_paged.cu @@ -726,6 +726,12 @@ __global__ void __launch_bounds__( first_n_tile = partition_id * tiles_per_partition; last_n_tile = min(num_n_tiles, first_n_tile + tiles_per_partition); } + if constexpr (!ANCHORED_SWA) { + // DFlash2 draft attention is non-causal with a short sliding window. Jump + // directly to the first tile that can intersect the window instead of + // visiting every historical tile only to reject it inside the loop. + first_n_tile = max(first_n_tile, min_key_pos / BLOCK_N); + } for (int block_n = first_n_tile; block_n < last_n_tile; ++block_n) { const int start_col = block_n * BLOCK_N; diff --git a/flash-attention-v100/kernel/fused_mha_forward_paged_full.cu b/flash-attention-v100/kernel/fused_mha_forward_paged_full.cu index bbcb8f8397..e9453037b5 100644 --- a/flash-attention-v100/kernel/fused_mha_forward_paged_full.cu +++ b/flash-attention-v100/kernel/fused_mha_forward_paged_full.cu @@ -16,589 +16,621 @@ using namespace nvcuda::wmma; #define WMMA_N 16 #define WMMA_K 16 -#define MAX_THREADS_PER_WARP 32 -#define MAX_THREADS_PER_SM 2048 -#define MAX_THREAD_BLOCK_SIZE 1024 +#define MAX_THREADS_PER_WARP 32 +#define MAX_THREADS_PER_SM 2048 +#define MAX_THREAD_BLOCK_SIZE 1024 #define MAX_THREAD_BLOCK_PER_SM 32 -#define MAX_WARPS_PER_SM 64 -#define MAX_SM_PER_GPU 80 -#define MAX_SMEM_PER_SM 98304 +#define MAX_WARPS_PER_SM 64 +#define MAX_SM_PER_GPU 80 +#define MAX_SMEM_PER_SM 98304 -#define WARP_ALLOC_GROUP 4 +#define WARP_ALLOC_GROUP 4 -#define MAX_REG_PER_UNIT 256 -#define MAX_REG_PER_THREAD 255 -#define MAX_REG_PER_BLOCK 65536 -#define MAX_REG_BUFFER 65536 +#define MAX_REG_PER_UNIT 256 +#define MAX_REG_PER_THREAD 255 +#define MAX_REG_PER_BLOCK 65536 +#define MAX_REG_BUFFER 65536 -#define BLOCK_M_16 16 -#define BLOCK_N_16 512 -#define WARPS_16 16 +#define BLOCK_M_16 16 +#define BLOCK_N_16 512 +#define WARPS_16 16 -#define BLOCK_M_32 32 -#define BLOCK_N_32 256 -#define WARPS_32 16 +#define BLOCK_M_32 32 +#define BLOCK_N_32 256 +#define WARPS_32 16 -#define BLOCK_M_64 64 -#define BLOCK_N_64 128 -#define WARPS_64 16 +#define BLOCK_M_64 64 +#define BLOCK_N_64 128 +#define WARPS_64 16 #define BLOCK_M_128 32 #define BLOCK_N_128 176 -#define WARPS_128 16 +#define WARPS_128 16 #define BLOCK_M_256 32 #define BLOCK_N_256 64 -#define WARPS_256 16 +#define WARPS_256 16 -template +template struct KernelConfig { - static constexpr int BLOCK_M = (D == 16) ? BLOCK_M_16 : (D == 32) ? BLOCK_M_32 : (D == 64) ? BLOCK_M_64 : (D == 128) ? BLOCK_M_128 : BLOCK_M_256; - static constexpr int BLOCK_N = (D == 16) ? BLOCK_N_16 : (D == 32) ? BLOCK_N_32 : (D == 64) ? BLOCK_N_64 : (D == 128) ? BLOCK_N_128 : BLOCK_N_256; - static constexpr int WARPS_PER_BLOCK = (D == 16) ? WARPS_16 : (D == 32) ? WARPS_32 : (D == 64) ? WARPS_64 : (D == 128) ? WARPS_128 : WARPS_256; - - static constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * MAX_THREADS_PER_WARP; - static constexpr int THREADS_PER_ROW = THREADS_PER_BLOCK / BLOCK_M; - static constexpr int PAD = (8 - (D % 32) + 32) % 32; - static constexpr int Q_STRIDE = D + PAD; - static constexpr int KV_STRIDE = D + PAD; - static constexpr int S_STRIDE = BLOCK_N + PAD; - static constexpr int O_STRIDE = D + PAD; - static constexpr int PER_UINT4 = 8; - - struct alignas(128) SmemLayout { - alignas(16) __half q [BLOCK_M * Q_STRIDE]; + static constexpr int BLOCK_M = (D == 16) ? BLOCK_M_16 + : (D == 32) ? BLOCK_M_32 + : (D == 64) ? BLOCK_M_64 + : (D == 128) ? BLOCK_M_128 + : BLOCK_M_256; + static constexpr int BLOCK_N = (D == 16) ? BLOCK_N_16 + : (D == 32) ? BLOCK_N_32 + : (D == 64) ? BLOCK_N_64 + : (D == 128) ? BLOCK_N_128 + : BLOCK_N_256; + static constexpr int WARPS_PER_BLOCK = (D == 16) ? WARPS_16 + : (D == 32) ? WARPS_32 + : (D == 64) ? WARPS_64 + : (D == 128) ? WARPS_128 + : WARPS_256; + + static constexpr int THREADS_PER_BLOCK = + WARPS_PER_BLOCK * MAX_THREADS_PER_WARP; + static constexpr int THREADS_PER_ROW = THREADS_PER_BLOCK / BLOCK_M; + static constexpr int PAD = (8 - (D % 32) + 32) % 32; + static constexpr int Q_STRIDE = D + PAD; + static constexpr int KV_STRIDE = D + PAD; + static constexpr int S_STRIDE = BLOCK_N + PAD; + static constexpr int O_STRIDE = D + PAD; + static constexpr int PER_UINT4 = 8; + + struct alignas(128) SmemLayout { + alignas(16) __half q[BLOCK_M * Q_STRIDE]; union { - alignas(16) __half k [BLOCK_N * KV_STRIDE]; - alignas(16) __half v [BLOCK_N * KV_STRIDE]; + alignas(16) __half k[BLOCK_N * KV_STRIDE]; + alignas(16) __half v[BLOCK_N * KV_STRIDE]; } reuse_kv; union { - alignas(16) float s [BLOCK_M * S_STRIDE]; - alignas(16) __half p [BLOCK_M * S_STRIDE]; + alignas(16) float s[BLOCK_M * S_STRIDE]; + alignas(16) __half p[BLOCK_M * S_STRIDE]; } reuse_sp; - alignas(16) float o [BLOCK_M * O_STRIDE]; - alignas(16) float row_max[BLOCK_M]; - alignas(16) float row_sum[BLOCK_M]; - }; + alignas(16) float o[BLOCK_M * O_STRIDE]; + alignas(16) float row_max[BLOCK_M]; + alignas(16) float row_sum[BLOCK_M]; + }; - static constexpr size_t TOTAL_SMEM = ((sizeof(SmemLayout) + 127) & ~size_t(127)); + static constexpr size_t TOTAL_SMEM = + ((sizeof(SmemLayout) + 127) & ~size_t(127)); }; -template +template __device__ __forceinline__ void init_smem(char* smem_raw) { - constexpr int N_U4 = Config::TOTAL_SMEM / 16; - constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; - const int tid = threadIdx.x; - - uint32_t addr = static_cast(__cvta_generic_to_shared(smem_raw)); - #pragma unroll 1 - for (int i = tid; i < N_U4; i += THREADS_PER_BLOCK) { - asm volatile("st.shared.v4.u32 [%0], {%1,%1,%1,%1};" - :: "r"(addr + (i << 4)), "r"(0) : "memory"); - } - __syncthreads(); + constexpr int N_U4 = Config::TOTAL_SMEM / 16; + constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; + const int tid = threadIdx.x; + + uint32_t addr = static_cast(__cvta_generic_to_shared(smem_raw)); +#pragma unroll 1 + for (int i = tid; i < N_U4; i += THREADS_PER_BLOCK) { + asm volatile("st.shared.v4.u32 [%0], {%1,%1,%1,%1};" ::"r"(addr + (i << 4)), + "r"(0) + : "memory"); + } + __syncthreads(); } -template +template __global__ void __launch_bounds__(KernelConfig::THREADS_PER_BLOCK, 2) -flash_attention_forward_kernel( - const __half* __restrict__ Q, - const __half* __restrict__ K, - const __half* __restrict__ V, - __half* __restrict__ Out, - float* __restrict__ softmax_lse, - const int B, - const int H, - const int M, - const int N, - const float softmax_scale -) { - using Config = KernelConfig; - constexpr int BLOCK_M = Config::BLOCK_M; - constexpr int BLOCK_N = Config::BLOCK_N; - constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; - constexpr int THREADS_PER_ROW = Config::THREADS_PER_ROW; - constexpr int WARPS_PER_BLOCK = Config::WARPS_PER_BLOCK; - constexpr int Q_STRIDE = Config::Q_STRIDE; - constexpr int KV_STRIDE = Config::KV_STRIDE; - constexpr int S_STRIDE = Config::S_STRIDE; - constexpr int O_STRIDE = Config::O_STRIDE; - constexpr int PER_UINT4 = Config::PER_UINT4; - const float NEG_INF = -1e30f; - - const int batch_head_id = blockIdx.z; - if (batch_head_id >= B * H) return; - - const int block_m = blockIdx.x; - const int start_row = block_m * BLOCK_M; - if (start_row >= M) return; - - int num_n_tiles = (N + BLOCK_N - 1) / BLOCK_N; - const int valid_q_rows = min(BLOCK_M, M - start_row); + flash_attention_forward_kernel(const __half* __restrict__ Q, + const __half* __restrict__ K, + const __half* __restrict__ V, + __half* __restrict__ Out, + float* __restrict__ softmax_lse, const int B, + const int H, const int M, const int N, + const float softmax_scale) { + using Config = KernelConfig; + constexpr int BLOCK_M = Config::BLOCK_M; + constexpr int BLOCK_N = Config::BLOCK_N; + constexpr int THREADS_PER_BLOCK = Config::THREADS_PER_BLOCK; + constexpr int THREADS_PER_ROW = Config::THREADS_PER_ROW; + constexpr int WARPS_PER_BLOCK = Config::WARPS_PER_BLOCK; + constexpr int Q_STRIDE = Config::Q_STRIDE; + constexpr int KV_STRIDE = Config::KV_STRIDE; + constexpr int S_STRIDE = Config::S_STRIDE; + constexpr int O_STRIDE = Config::O_STRIDE; + constexpr int PER_UINT4 = Config::PER_UINT4; + const float NEG_INF = -1e30f; + + const int batch_head_id = blockIdx.z; + if (batch_head_id >= B * H) return; + + const int block_m = blockIdx.x; + const int start_row = block_m * BLOCK_M; + if (start_row >= M) return; + + int num_n_tiles = (N + BLOCK_N - 1) / BLOCK_N; + const int valid_q_rows = min(BLOCK_M, M - start_row); + + if constexpr (IS_CAUSAL) { + const int max_key_pos = start_row + valid_q_rows - 1; + if (max_key_pos < 0) { + num_n_tiles = 0; + } else { + num_n_tiles = min(num_n_tiles, (max_key_pos + BLOCK_N + 0) / BLOCK_N); + } + } + + const int tid = threadIdx.x; + const int warp_id = tid / MAX_THREADS_PER_WARP; + const int lane_id = tid % MAX_THREADS_PER_WARP; + + const __half* q_ptr = Q + (size_t)batch_head_id * M * D + start_row * D; + const __half* k_ptr = K + (size_t)batch_head_id * N * D; + const __half* v_ptr = V + (size_t)batch_head_id * N * D; + __half* out_ptr = Out + (size_t)batch_head_id * M * D + start_row * D; + float* softmax_lse_ptr = softmax_lse + (size_t)batch_head_id * M + start_row; + + extern __shared__ char smem_raw[]; + init_smem(smem_raw); + auto& smem = *reinterpret_cast(smem_raw); + + __half* sQ = smem.q; + __half* sK = smem.reuse_kv.k; + __half* sV = smem.reuse_kv.v; + float* sS = smem.reuse_sp.s; + __half* sP = smem.reuse_sp.p; + float* sO = smem.o; + float* sRowMax = smem.row_max; + float* sRowSum = smem.row_sum; + + const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; + const int q_stride_uint4 = (Q_STRIDE + PER_UINT4 - 1) / PER_UINT4; + const int kv_stride_uint4 = (KV_STRIDE + PER_UINT4 - 1) / PER_UINT4; + + if (tid < BLOCK_M) { + sRowMax[tid] = NEG_INF; + } + + const uint4* q_vec = reinterpret_cast(q_ptr); + uint4* sQ_vec = reinterpret_cast(sQ); + +#pragma unroll 4 + for (int idx = tid; idx < (valid_q_rows * d_stride_uint4); + idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; + uint4 q_val = make_uint4(0, 0, 0, 0); + if (row < valid_q_rows && vec_col < d_stride_uint4) { + q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); + } + sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + } + __syncthreads(); + + for (int block_n = 0; block_n < num_n_tiles; ++block_n) { + const int start_col = block_n * BLOCK_N; + if (start_col >= N) break; + const int valid_k_rows = min(BLOCK_N, N - start_col); if constexpr (IS_CAUSAL) { - const int max_key_pos = start_row + valid_q_rows - 1; - if (max_key_pos < 0) { - num_n_tiles = 0; - } else { - num_n_tiles = min(num_n_tiles, (max_key_pos + BLOCK_N + 0) / BLOCK_N); - } + if (start_col >= start_row + valid_q_rows) { + continue; + } } - const int tid = threadIdx.x; - const int warp_id = tid / MAX_THREADS_PER_WARP; - const int lane_id = tid % MAX_THREADS_PER_WARP; - - const __half* q_ptr = Q + (size_t)batch_head_id * M * D + start_row * D; - const __half* k_ptr = K + (size_t)batch_head_id * N * D; - const __half* v_ptr = V + (size_t)batch_head_id * N * D; - __half* out_ptr = Out + (size_t)batch_head_id * M * D + start_row * D; - float* softmax_lse_ptr = softmax_lse + (size_t)batch_head_id * M + start_row; - - extern __shared__ char smem_raw[]; - init_smem(smem_raw); - auto& smem = *reinterpret_cast(smem_raw); - - __half* sQ = smem.q; - __half* sK = smem.reuse_kv.k; - __half* sV = smem.reuse_kv.v; - float* sS = smem.reuse_sp.s; - __half* sP = smem.reuse_sp.p; - float* sO = smem.o; - float* sRowMax = smem.row_max; - float* sRowSum = smem.row_sum; - - const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; - const int q_stride_uint4 = (Q_STRIDE + PER_UINT4 - 1) / PER_UINT4; - const int kv_stride_uint4 = (KV_STRIDE + PER_UINT4 - 1) / PER_UINT4; - - if (tid < BLOCK_M) { - sRowMax[tid] = NEG_INF; - } + const uint4* k_vec = reinterpret_cast(k_ptr + start_col * D); + uint4* sK_vec = reinterpret_cast(sK); - const uint4* q_vec = reinterpret_cast(q_ptr); - uint4* sQ_vec = reinterpret_cast(sQ); +#pragma unroll 2 + for (int idx = tid; idx < (valid_k_rows * d_stride_uint4); + idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; - #pragma unroll 4 - for (int idx = tid; idx < (valid_q_rows * d_stride_uint4); idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - uint4 q_val = make_uint4(0, 0, 0, 0); - if (row < valid_q_rows && vec_col < d_stride_uint4) { - q_val = __ldg(&q_vec[row * d_stride_uint4 + vec_col]); - } - sQ_vec[row * q_stride_uint4 + vec_col] = q_val; + uint4 k_val = make_uint4(0, 0, 0, 0); + if (row < valid_k_rows && vec_col < d_stride_uint4) { + k_val = __ldg(&k_vec[row * d_stride_uint4 + vec_col]); + } + sK_vec[row * kv_stride_uint4 + vec_col] = k_val; } __syncthreads(); - for (int block_n = 0; block_n < num_n_tiles; ++block_n) { - const int start_col = block_n * BLOCK_N; - if (start_col >= N) break; - const int valid_k_rows = min(BLOCK_N, N - start_col); - - if constexpr (IS_CAUSAL) { - if (start_col >= start_row + valid_q_rows) { continue; } + const int num_tiles_m_qk = (BLOCK_M + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_qk = (BLOCK_N + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_qk = (D + WMMA_K - 1) / WMMA_K; + const int total_tiles_qk = num_tiles_m_qk * num_tiles_n_qk; + const int tiles_per_warp_qk = + (total_tiles_qk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + const unsigned row_causal = (lane_id & 0b1) + ((lane_id >> 2) & 0b1) * 8 + + ((lane_id >> 4) & 0b1) * 4; + const unsigned col_causal = + ((lane_id >> 1) & 0b1) * 2 + ((lane_id >> 3) & 0b1) * 8; + + for (int tile_idx = 0; tile_idx < tiles_per_warp_qk; ++tile_idx) { + const int global_tile_idx = warp_id * tiles_per_warp_qk + tile_idx; + if (global_tile_idx >= total_tiles_qk) break; + + const int tile_m_idx = global_tile_idx / num_tiles_n_qk; + const int tile_n_idx = global_tile_idx % num_tiles_n_qk; + + const int tile_m = tile_m_idx * WMMA_M; + const int tile_n = tile_n_idx * WMMA_N; + + if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; + + fragment a_frag; + fragment b_frag; + fragment acc_frag; + fill_fragment(acc_frag, 0.0f); + +#pragma unroll + for (int k_tile = 0; k_tile < num_tiles_k_qk; ++k_tile) { + const int k_offset = k_tile * WMMA_K; + if (k_offset >= D) break; + + load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, Q_STRIDE); + load_matrix_sync(b_frag, sK + tile_n * KV_STRIDE + k_offset, KV_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + if constexpr (IS_CAUSAL) { +#pragma unroll + for (int i = 0; i < acc_frag.num_elements; ++i) { + const unsigned col = col_causal + (i & 0b1) + ((i >> 2) & 0b1) * 4; + const unsigned row = row_causal + ((i >> 1) & 0b1) * 2; + + const int global_m = start_row + tile_m + row; + const int global_n = start_col + tile_n + col; + + const bool is_valid = (global_m < start_row + valid_q_rows) && + (global_n < start_col + valid_k_rows); + + acc_frag.x[i] = + is_valid ? ((global_n > global_m) ? NEG_INF + : acc_frag.x[i] * softmax_scale) + : NEG_INF; } + } else { +#pragma unroll + for (int i = 0; i < acc_frag.num_elements; ++i) { + acc_frag.x[i] *= softmax_scale; + } + } + store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, + mem_row_major); + } + __syncthreads(); - const uint4* k_vec = reinterpret_cast(k_ptr + start_col * D); - uint4* sK_vec = reinterpret_cast(sK); - - #pragma unroll 2 - for (int idx = tid; idx < (valid_k_rows * d_stride_uint4); idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - - uint4 k_val = make_uint4(0, 0, 0, 0); - if (row < valid_k_rows && vec_col < d_stride_uint4) { - k_val = __ldg(&k_vec[row * d_stride_uint4 + vec_col]); - } - sK_vec[row * kv_stride_uint4 + vec_col] = k_val; + if (tid < valid_q_rows * THREADS_PER_ROW) { + const int row = tid / THREADS_PER_ROW; + const int thread_in_row = tid % THREADS_PER_ROW; + const unsigned mask = + (valid_q_rows == BLOCK_M) ? 0xFFFFFFFFU : __activemask(); + const int row_leader = __ffs(mask) - 1; + + float* sS_row_f = sS + row * S_STRIDE; + __half* sP_row_h = sP + row * S_STRIDE; + + const int vec_cols = valid_k_rows >> 2; + const int vecs_per_thread = + (vec_cols + THREADS_PER_ROW - 1) / THREADS_PER_ROW; + const int tail_start = vec_cols << 2; + + float thread_max = NEG_INF; + float4* sS_vec4 = reinterpret_cast(sS_row_f); + +#pragma unroll 4 + for (int j = 0; j < vecs_per_thread; ++j) { + int vc = thread_in_row + j * THREADS_PER_ROW; + if (vc < vec_cols) { + float4 v4 = sS_vec4[vc]; + thread_max = + fmaxf(thread_max, fmaxf(fmaxf(v4.x, v4.y), fmaxf(v4.z, v4.w))); } - __syncthreads(); - - const int num_tiles_m_qk = (BLOCK_M + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_qk = (BLOCK_N + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_qk = (D + WMMA_K - 1) / WMMA_K; - const int total_tiles_qk = num_tiles_m_qk * num_tiles_n_qk; - const int tiles_per_warp_qk = (total_tiles_qk + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - const unsigned row_causal = (lane_id & 0b1) + ((lane_id >> 2) & 0b1) * 8 + ((lane_id >> 4) & 0b1) * 4; - const unsigned col_causal = ((lane_id >> 1) & 0b1) * 2 + ((lane_id >> 3) & 0b1) * 8; - - for (int tile_idx = 0; tile_idx < tiles_per_warp_qk; ++tile_idx) { - const int global_tile_idx = warp_id * tiles_per_warp_qk + tile_idx; - if (global_tile_idx >= total_tiles_qk) break; - - const int tile_m_idx = global_tile_idx / num_tiles_n_qk; - const int tile_n_idx = global_tile_idx % num_tiles_n_qk; - - const int tile_m = tile_m_idx * WMMA_M; - const int tile_n = tile_n_idx * WMMA_N; - - if (tile_m >= valid_q_rows || tile_n >= valid_k_rows) continue; - - fragment a_frag; - fragment b_frag; - fragment acc_frag; - fill_fragment(acc_frag, 0.0f); - - #pragma unroll - for (int k_tile = 0; k_tile < num_tiles_k_qk; ++k_tile) { - const int k_offset = k_tile * WMMA_K; - if (k_offset >= D) break; - - load_matrix_sync(a_frag, sQ + tile_m * Q_STRIDE + k_offset, Q_STRIDE); - load_matrix_sync(b_frag, sK + tile_n * KV_STRIDE + k_offset, KV_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - - if constexpr (IS_CAUSAL) { - #pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - const unsigned col = col_causal + (i & 0b1) + ((i >> 2) & 0b1) * 4; - const unsigned row = row_causal + ((i >> 1) & 0b1) * 2; - - const int global_m = start_row + tile_m + row; - const int global_n = start_col + tile_n + col; - - const bool is_valid = (global_m < start_row + valid_q_rows) && - (global_n < start_col + valid_k_rows); - - acc_frag.x[i] = is_valid - ? ((global_n > global_m) ? NEG_INF : acc_frag.x[i] * softmax_scale) - : NEG_INF; - } - } else { - #pragma unroll - for (int i = 0; i < acc_frag.num_elements; ++i) { - acc_frag.x[i] *= softmax_scale; - } - } - store_matrix_sync(sS + tile_m * S_STRIDE + tile_n, acc_frag, S_STRIDE, mem_row_major); + } + +#pragma unroll 4 + for (int c = tail_start + thread_in_row; c < valid_k_rows; + c += THREADS_PER_ROW) { + thread_max = fmaxf(thread_max, sS_row_f[c]); + } + +#pragma unroll + for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) + thread_max = fmaxf( + thread_max, __shfl_down_sync(mask, thread_max, o, THREADS_PER_ROW)); + + const float row_max = + __shfl_sync(mask, thread_max, row_leader, THREADS_PER_ROW); + const float old_max = sRowMax[row]; + const float new_max = fmaxf(old_max, row_max); + const float exp_diff = __expf(old_max - new_max); + + float thread_sum = 0.0f; + __half2 half_buffer[20]; + int vc_base = thread_in_row; + int h2_idx = 0; + int tail_col = -1; + __half tail_value = __float2half(0.f); + +#pragma unroll 4 + for (int j = 0; j < vecs_per_thread; ++j, vc_base += THREADS_PER_ROW) { + if (vc_base < vec_cols) { + float4 v4 = sS_vec4[vc_base]; + + float e0 = __expf(fmaxf(v4.x - new_max, -80.0f)); + float e1 = __expf(fmaxf(v4.y - new_max, -80.0f)); + float e2 = __expf(fmaxf(v4.z - new_max, -80.0f)); + float e3 = __expf(fmaxf(v4.w - new_max, -80.0f)); + + thread_sum += (e0 + e1) + (e2 + e3); + + half_buffer[h2_idx++] = __float22half2_rn(make_float2(e0, e1)); + half_buffer[h2_idx++] = __float22half2_rn(make_float2(e2, e3)); } - __syncthreads(); - - if (tid < valid_q_rows * THREADS_PER_ROW) { - const int row = tid / THREADS_PER_ROW; - const int thread_in_row = tid % THREADS_PER_ROW; - const unsigned mask = (valid_q_rows == BLOCK_M) ? 0xFFFFFFFFU : __activemask(); - const int row_leader = __ffs(mask) - 1; - - float* sS_row_f = sS + row * S_STRIDE; - __half* sP_row_h = sP + row * S_STRIDE; - - const int vec_cols = valid_k_rows >> 2; - const int vecs_per_thread = (vec_cols + THREADS_PER_ROW - 1) / THREADS_PER_ROW; - const int tail_start = vec_cols << 2; - - float thread_max = NEG_INF; - float4* sS_vec4 = reinterpret_cast(sS_row_f); - - #pragma unroll 4 - for (int j = 0; j < vecs_per_thread; ++j) { - int vc = thread_in_row + j * THREADS_PER_ROW; - if (vc < vec_cols) { - float4 v4 = sS_vec4[vc]; - thread_max = fmaxf(thread_max, fmaxf(fmaxf(v4.x, v4.y), fmaxf(v4.z, v4.w))); - } - } - - #pragma unroll 4 - for (int c = tail_start + thread_in_row; c < valid_k_rows; c += THREADS_PER_ROW) { - thread_max = fmaxf(thread_max, sS_row_f[c]); - } - - #pragma unroll - for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) - thread_max = fmaxf(thread_max, __shfl_down_sync(mask, thread_max, o, THREADS_PER_ROW)); - - const float row_max = __shfl_sync(mask, thread_max, row_leader, THREADS_PER_ROW); - const float old_max = sRowMax[row]; - const float new_max = fmaxf(old_max, row_max); - const float exp_diff = __expf(old_max - new_max); - - float thread_sum = 0.0f; - __half2 half_buffer[20]; - int vc_base = thread_in_row; - int h2_idx = 0; - int tail_col = -1; - __half tail_value = __float2half(0.f); - - #pragma unroll 4 - for (int j = 0; j < vecs_per_thread; ++j, vc_base += THREADS_PER_ROW) { - if (vc_base < vec_cols) { - float4 v4 = sS_vec4[vc_base]; - - float e0 = __expf(fmaxf(v4.x - new_max, -80.0f)); - float e1 = __expf(fmaxf(v4.y - new_max, -80.0f)); - float e2 = __expf(fmaxf(v4.z - new_max, -80.0f)); - float e3 = __expf(fmaxf(v4.w - new_max, -80.0f)); - - thread_sum += (e0 + e1) + (e2 + e3); - - half_buffer[h2_idx++] = __float22half2_rn(make_float2(e0, e1)); - half_buffer[h2_idx++] = __float22half2_rn(make_float2(e2, e3)); - } - } - - #pragma unroll 4 - for (int c = tail_start + thread_in_row; c < valid_k_rows; c += THREADS_PER_ROW) { - float v = sS_row_f[c]; - float e = __expf(fmaxf(v - new_max, -80.0f)); - thread_sum += e; - tail_col = c; - tail_value = __float2half_rn(e); - } - - #pragma unroll - for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) - thread_sum += __shfl_down_sync(mask, thread_sum, o, THREADS_PER_ROW); - - float row_sum = __shfl_sync(mask, thread_sum, row_leader, THREADS_PER_ROW); - - if (thread_in_row == 0) { - sRowSum[row] = exp_diff * sRowSum[row] + row_sum; - sRowMax[row] = new_max; - } - - h2_idx = 0; - vc_base = thread_in_row; - __half2* sP_half2 = reinterpret_cast<__half2*>(sP_row_h); - - #pragma unroll 4 - for (int j = 0; j < vecs_per_thread; ++j, vc_base += THREADS_PER_ROW) { - if (vc_base < vec_cols) { - int base_offset = vc_base * 2; - - sP_half2[base_offset] = half_buffer[h2_idx++]; - sP_half2[base_offset + 1] = half_buffer[h2_idx++]; - } - } - - if (tail_col >= 0) { - sP_row_h[tail_col] = tail_value; - } - - #pragma unroll 4 - for (int c = tail_start + thread_in_row; c < BLOCK_N; c += THREADS_PER_ROW) { - if (c >= valid_k_rows) { - sP_row_h[c] = __float2half(0.f); - } - } - - if (block_n > 0) { - float* sO_row = sO + row * O_STRIDE; - float4* sO_vec = reinterpret_cast(sO_row); - const int o_vec_count = (O_STRIDE + 3) >> 2; - float scale = exp_diff; - - #pragma unroll 4 - for (int ov = thread_in_row; ov < o_vec_count; ov += THREADS_PER_ROW) { - float4 v = sO_vec[ov]; - v.x *= scale; - v.y *= scale; - v.z *= scale; - v.w *= scale; - - sO_vec[ov] = v; - } - } + } + +#pragma unroll 4 + for (int c = tail_start + thread_in_row; c < valid_k_rows; + c += THREADS_PER_ROW) { + float v = sS_row_f[c]; + float e = __expf(fmaxf(v - new_max, -80.0f)); + thread_sum += e; + tail_col = c; + tail_value = __float2half_rn(e); + } + +#pragma unroll + for (int o = THREADS_PER_ROW / 2; o > 0; o >>= 1) + thread_sum += __shfl_down_sync(mask, thread_sum, o, THREADS_PER_ROW); + + float row_sum = + __shfl_sync(mask, thread_sum, row_leader, THREADS_PER_ROW); + + if (thread_in_row == 0) { + sRowSum[row] = exp_diff * sRowSum[row] + row_sum; + sRowMax[row] = new_max; + } + + h2_idx = 0; + vc_base = thread_in_row; + __half2* sP_half2 = reinterpret_cast<__half2*>(sP_row_h); + +#pragma unroll 4 + for (int j = 0; j < vecs_per_thread; ++j, vc_base += THREADS_PER_ROW) { + if (vc_base < vec_cols) { + int base_offset = vc_base * 2; + + sP_half2[base_offset] = half_buffer[h2_idx++]; + sP_half2[base_offset + 1] = half_buffer[h2_idx++]; } - __syncthreads(); + } - const uint4* v_vec = reinterpret_cast(v_ptr + start_col * D); - uint4* sV_vec = reinterpret_cast(sV); + if (tail_col >= 0) { + sP_row_h[tail_col] = tail_value; + } - #pragma unroll 2 - for (int idx = tid; idx < (valid_k_rows * d_stride_uint4); idx += THREADS_PER_BLOCK) { - const int row = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; - - uint4 v_val = make_uint4(0, 0, 0, 0); - if (row < valid_k_rows && vec_col < d_stride_uint4) { - v_val = __ldg(&v_vec[row * d_stride_uint4 + vec_col]); - } - sV_vec[row * kv_stride_uint4 + vec_col] = v_val; +#pragma unroll 4 + for (int c = tail_start + thread_in_row; c < BLOCK_N; + c += THREADS_PER_ROW) { + if (c >= valid_k_rows) { + sP_row_h[c] = __float2half(0.f); + } + } + + if (block_n > 0) { + float* sO_row = sO + row * O_STRIDE; + float4* sO_vec = reinterpret_cast(sO_row); + const int o_vec_count = (O_STRIDE + 3) >> 2; + float scale = exp_diff; + +#pragma unroll 4 + for (int ov = thread_in_row; ov < o_vec_count; ov += THREADS_PER_ROW) { + float4 v = sO_vec[ov]; + v.x *= scale; + v.y *= scale; + v.z *= scale; + v.w *= scale; + + sO_vec[ov] = v; } - __syncthreads(); + } + } + __syncthreads(); - const int num_tiles_m_pv = (BLOCK_M + WMMA_M - 1) / WMMA_M; - const int num_tiles_n_pv = (D + WMMA_N - 1) / WMMA_N; - const int num_tiles_k_pv = (BLOCK_N + WMMA_K - 1) / WMMA_K; - const int total_tiles_pv = num_tiles_m_pv * num_tiles_n_pv; - const int tiles_per_warp_pv = (total_tiles_pv + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + const uint4* v_vec = reinterpret_cast(v_ptr + start_col * D); + uint4* sV_vec = reinterpret_cast(sV); - for (int tile_idx = 0; tile_idx < tiles_per_warp_pv; ++tile_idx) { - const int global_tile_idx = warp_id * tiles_per_warp_pv + tile_idx; - if (global_tile_idx >= total_tiles_pv) break; +#pragma unroll 2 + for (int idx = tid; idx < (valid_k_rows * d_stride_uint4); + idx += THREADS_PER_BLOCK) { + const int row = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; - const int tile_m_idx = global_tile_idx / num_tiles_n_pv; - const int tile_d_idx = global_tile_idx % num_tiles_n_pv; + uint4 v_val = make_uint4(0, 0, 0, 0); + if (row < valid_k_rows && vec_col < d_stride_uint4) { + v_val = __ldg(&v_vec[row * d_stride_uint4 + vec_col]); + } + sV_vec[row * kv_stride_uint4 + vec_col] = v_val; + } + __syncthreads(); - const int tile_m = tile_m_idx * WMMA_M; - const int tile_d = tile_d_idx * WMMA_N; + const int num_tiles_m_pv = (BLOCK_M + WMMA_M - 1) / WMMA_M; + const int num_tiles_n_pv = (D + WMMA_N - 1) / WMMA_N; + const int num_tiles_k_pv = (BLOCK_N + WMMA_K - 1) / WMMA_K; + const int total_tiles_pv = num_tiles_m_pv * num_tiles_n_pv; + const int tiles_per_warp_pv = + (total_tiles_pv + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; - if (tile_m >= valid_q_rows) continue; + for (int tile_idx = 0; tile_idx < tiles_per_warp_pv; ++tile_idx) { + const int global_tile_idx = warp_id * tiles_per_warp_pv + tile_idx; + if (global_tile_idx >= total_tiles_pv) break; - fragment a_frag; - fragment b_frag; - fragment acc_frag; + const int tile_m_idx = global_tile_idx / num_tiles_n_pv; + const int tile_d_idx = global_tile_idx % num_tiles_n_pv; - load_matrix_sync(acc_frag, sO + tile_m * O_STRIDE + tile_d, O_STRIDE, mem_row_major); + const int tile_m = tile_m_idx * WMMA_M; + const int tile_d = tile_d_idx * WMMA_N; - #pragma unroll - for (int tile_k = 0; tile_k < num_tiles_k_pv; ++tile_k) { - const int k_offset = tile_k * WMMA_K; - if (k_offset >= valid_k_rows) break; + if (tile_m >= valid_q_rows) continue; - load_matrix_sync(a_frag, sP + tile_m * S_STRIDE + k_offset, S_STRIDE); - load_matrix_sync(b_frag, sV + k_offset * KV_STRIDE + tile_d, KV_STRIDE); - mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - store_matrix_sync(sO + tile_m * O_STRIDE + tile_d, acc_frag, O_STRIDE, mem_row_major); - } - __syncthreads(); - } + fragment a_frag; + fragment b_frag; + fragment acc_frag; - const int total_fp16_x4 = (valid_q_rows * D) / 4; - - for (int i = tid; i < total_fp16_x4; i += THREADS_PER_BLOCK) { - const int row = i / (D / 4); - const int col = (i % (D / 4)) * 4; - - const float sum_clamped = fmaxf(sRowSum[row], 1e-24f); - const float inv_sum = 1.0f / sum_clamped; - const float* sO_row = sO + row * O_STRIDE; - - const __half h0 = __float2half_rn(sO_row[col + 0] * inv_sum); - const __half h1 = __float2half_rn(sO_row[col + 1] * inv_sum); - const __half h2 = __float2half_rn(sO_row[col + 2] * inv_sum); - const __half h3 = __float2half_rn(sO_row[col + 3] * inv_sum); - - asm volatile( - "st.global.v4.u16 [%0], {%1, %2, %3, %4};" - : - : "l"(out_ptr + row * D + col), - "h"(__half_as_ushort(h0)), - "h"(__half_as_ushort(h1)), - "h"(__half_as_ushort(h2)), - "h"(__half_as_ushort(h3)) - : "memory" - ); - } + load_matrix_sync(acc_frag, sO + tile_m * O_STRIDE + tile_d, O_STRIDE, + mem_row_major); + +#pragma unroll + for (int tile_k = 0; tile_k < num_tiles_k_pv; ++tile_k) { + const int k_offset = tile_k * WMMA_K; + if (k_offset >= valid_k_rows) break; - if (tid < valid_q_rows) { - const float sum = fmaxf(sRowSum[tid], 1e-24f); - softmax_lse_ptr[tid] = sRowMax[tid] + logf(sum); + load_matrix_sync(a_frag, sP + tile_m * S_STRIDE + k_offset, S_STRIDE); + load_matrix_sync(b_frag, sV + k_offset * KV_STRIDE + tile_d, KV_STRIDE); + mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + store_matrix_sync(sO + tile_m * O_STRIDE + tile_d, acc_frag, O_STRIDE, + mem_row_major); } + __syncthreads(); + } + + const int total_fp16_x4 = (valid_q_rows * D) / 4; + + for (int i = tid; i < total_fp16_x4; i += THREADS_PER_BLOCK) { + const int row = i / (D / 4); + const int col = (i % (D / 4)) * 4; + + const float sum_clamped = fmaxf(sRowSum[row], 1e-24f); + const float inv_sum = 1.0f / sum_clamped; + const float* sO_row = sO + row * O_STRIDE; + + const __half h0 = __float2half_rn(sO_row[col + 0] * inv_sum); + const __half h1 = __float2half_rn(sO_row[col + 1] * inv_sum); + const __half h2 = __float2half_rn(sO_row[col + 2] * inv_sum); + const __half h3 = __float2half_rn(sO_row[col + 3] * inv_sum); + + asm volatile("st.global.v4.u16 [%0], {%1, %2, %3, %4};" + : + : "l"(out_ptr + row * D + col), "h"(__half_as_ushort(h0)), + "h"(__half_as_ushort(h1)), "h"(__half_as_ushort(h2)), + "h"(__half_as_ushort(h3)) + : "memory"); + } + + if (tid < valid_q_rows) { + const float sum = fmaxf(sRowSum[tid], 1e-24f); + softmax_lse_ptr[tid] = sRowMax[tid] + logf(sum); + } } -template +template void launcher_flash_attention_forward( - const torch::Tensor& Q, - const torch::Tensor& K, - const torch::Tensor& V, - torch::Tensor& Out, - torch::Tensor& softmax_lse, - float softmax_scale, - bool is_causal, - cudaStream_t stream -) { - using Config = KernelConfig; - - const int B = Q.size(0); - const int H = Q.size(1); - const int M = Q.size(2); - const int N = K.size(2); - - const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; - const dim3 grid(grid_x, 1, B * H); - const dim3 block(Config::THREADS_PER_BLOCK); - const size_t smem = Config::TOTAL_SMEM; - - TORCH_CHECK(smem <= MAX_SMEM_PER_SM, "Shared memory exceeds 96KB: ", smem, " bytes"); - - auto kernel = is_causal ? - (void*)flash_attention_forward_kernel : - (void*)flash_attention_forward_kernel; - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - - if (is_causal) { - flash_attention_forward_kernel<<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K.data_ptr()), - reinterpret_cast(V.data_ptr()), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), - B, H, M, N, softmax_scale - ); - } else { - flash_attention_forward_kernel<<>>( - reinterpret_cast(Q.data_ptr()), - reinterpret_cast(K.data_ptr()), - reinterpret_cast(V.data_ptr()), - reinterpret_cast<__half*>(Out.data_ptr()), - softmax_lse.data_ptr(), - B, H, M, N, softmax_scale - ); - } + const torch::Tensor& Q, const torch::Tensor& K, const torch::Tensor& V, + torch::Tensor& Out, torch::Tensor& softmax_lse, float softmax_scale, + bool is_causal, cudaStream_t stream) { + using Config = KernelConfig; + + const int B = Q.size(0); + const int H = Q.size(1); + const int M = Q.size(2); + const int N = K.size(2); + + const int grid_x = (M + Config::BLOCK_M - 1) / Config::BLOCK_M; + const dim3 grid(grid_x, 1, B * H); + const dim3 block(Config::THREADS_PER_BLOCK); + const size_t smem = Config::TOTAL_SMEM; + + TORCH_CHECK(smem <= MAX_SMEM_PER_SM, "Shared memory exceeds 96KB: ", smem, + " bytes"); + + auto kernel = is_causal ? (void*)flash_attention_forward_kernel + : (void*)flash_attention_forward_kernel; + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); + + if (is_causal) { + flash_attention_forward_kernel<<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K.data_ptr()), + reinterpret_cast(V.data_ptr()), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), B, H, M, N, softmax_scale); + } else { + flash_attention_forward_kernel<<>>( + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K.data_ptr()), + reinterpret_cast(V.data_ptr()), + reinterpret_cast<__half*>(Out.data_ptr()), + softmax_lse.data_ptr(), B, H, M, N, softmax_scale); + } } std::vector flash_attention_forward( - at::Tensor& q, - const at::Tensor& k, - const at::Tensor& v, - std::optional& out_, - std::optional& alibi_slopes_, - const float p_dropout, - const float softmax_scale, - bool is_causal, - int window_size_left, - int window_size_right, - const float softcap, - const bool return_softmax, - std::optional gen_ -) { - - TORCH_CHECK(!alibi_slopes_.has_value(), "alibi_slopes not supported"); - TORCH_CHECK(p_dropout == 0.f, "dropout not supported"); - TORCH_CHECK(window_size_left == -1, "window_size_left not supported"); - TORCH_CHECK(window_size_right == -1 || (is_causal && window_size_right == 0), "window not supported"); - TORCH_CHECK(softcap == 0.f, "softcap not supported"); - TORCH_CHECK(!return_softmax, "return_softmax not supported"); - TORCH_CHECK(!gen_.has_value(), "Generator not supported"); - - TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); - TORCH_CHECK(k.dtype() == torch::kFloat16, "k must be fp16"); - TORCH_CHECK(v.dtype() == torch::kFloat16, "v must be fp16"); - TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "Tensors must be on CUDA"); - TORCH_CHECK(q.stride(-1) == 1 && k.stride(-1) == 1 && v.stride(-1) == 1, "Last dim must be contiguous"); - - const auto sizes = q.sizes(); - const int B = sizes[0], H = sizes[1], M = sizes[2], D = sizes[3]; - const int N = k.size(2); - TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, "D must be even, <=256, multiple of 8"); - - at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); - TORCH_CHECK(out_fp16.dtype() == torch::kFloat16, "out must be fp16"); - auto softmax_lse = torch::zeros({B, H, M}, torch::dtype(torch::kFloat32).device(q.device())); - TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, "softmax_lse must be fp32"); - - auto stream = at::cuda::getCurrentCUDAStream().stream(); - auto props = at::cuda::getCurrentDeviceProperties(); - bool sm70 = props->major == 7 && props->minor == 0; - TORCH_CHECK(sm70, "Kernel supports only Volta GPUs."); - - switch (D) { - case 16: launcher_flash_attention_forward<16>(q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, stream); break; - case 32: launcher_flash_attention_forward<32>(q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, stream); break; - case 64: launcher_flash_attention_forward<64>(q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, stream); break; - case 128: launcher_flash_attention_forward<128>(q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, stream); break; - case 256: launcher_flash_attention_forward<256>(q, k, v, out_fp16, softmax_lse, softmax_scale, is_causal, stream); break; - default: TORCH_CHECK(false, "Unsupported D: ", D); - } - - auto p = torch::zeros({0}, q.options()); - auto rng_state = torch::zeros({2}, torch::dtype(torch::kInt64).device(q.device())); - return {out_fp16, softmax_lse, p, rng_state}; + at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out_, std::optional& alibi_slopes_, + const float p_dropout, const float softmax_scale, bool is_causal, + int window_size_left, int window_size_right, const float softcap, + const bool return_softmax, std::optional gen_) { + TORCH_CHECK(!alibi_slopes_.has_value(), "alibi_slopes not supported"); + TORCH_CHECK(p_dropout == 0.f, "dropout not supported"); + TORCH_CHECK(window_size_left == -1, "window_size_left not supported"); + TORCH_CHECK(window_size_right == -1 || (is_causal && window_size_right == 0), + "window not supported"); + TORCH_CHECK(softcap == 0.f, "softcap not supported"); + TORCH_CHECK(!return_softmax, "return_softmax not supported"); + TORCH_CHECK(!gen_.has_value(), "Generator not supported"); + + TORCH_CHECK(q.dtype() == torch::kFloat16, "q must be fp16"); + TORCH_CHECK(k.dtype() == torch::kFloat16, "k must be fp16"); + TORCH_CHECK(v.dtype() == torch::kFloat16, "v must be fp16"); + TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), + "Tensors must be on CUDA"); + TORCH_CHECK(q.stride(-1) == 1 && k.stride(-1) == 1 && v.stride(-1) == 1, + "Last dim must be contiguous"); + + const auto sizes = q.sizes(); + const int B = sizes[0], H = sizes[1], M = sizes[2], D = sizes[3]; + const int N = k.size(2); + TORCH_CHECK(D <= 256 && D % 8 == 0 && D % 2 == 0, + "D must be even, <=256, multiple of 8"); + + at::Tensor out_fp16 = out_.has_value() ? out_.value() : torch::zeros_like(q); + TORCH_CHECK(out_fp16.dtype() == torch::kFloat16, "out must be fp16"); + auto softmax_lse = + torch::zeros({B, H, M}, torch::dtype(torch::kFloat32).device(q.device())); + TORCH_CHECK(softmax_lse.dtype() == torch::kFloat32, + "softmax_lse must be fp32"); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + auto props = at::cuda::getCurrentDeviceProperties(); + bool sm70 = props->major == 7 && props->minor == 0; + TORCH_CHECK(sm70, "Kernel supports only Volta GPUs."); + + switch (D) { + case 16: + launcher_flash_attention_forward<16>(q, k, v, out_fp16, softmax_lse, + softmax_scale, is_causal, stream); + break; + case 32: + launcher_flash_attention_forward<32>(q, k, v, out_fp16, softmax_lse, + softmax_scale, is_causal, stream); + break; + case 64: + launcher_flash_attention_forward<64>(q, k, v, out_fp16, softmax_lse, + softmax_scale, is_causal, stream); + break; + case 128: + launcher_flash_attention_forward<128>(q, k, v, out_fp16, softmax_lse, + softmax_scale, is_causal, stream); + break; + case 256: + launcher_flash_attention_forward<256>(q, k, v, out_fp16, softmax_lse, + softmax_scale, is_causal, stream); + break; + default: + TORCH_CHECK(false, "Unsupported D: ", D); + } + + auto p = torch::zeros({0}, q.options()); + auto rng_state = + torch::zeros({2}, torch::dtype(torch::kInt64).device(q.device())); + return {out_fp16, softmax_lse, p, rng_state}; } diff --git a/flash-attention-v100/kernel/paged_kv_utils.cuh b/flash-attention-v100/kernel/paged_kv_utils.cuh index 18c8274d2e..bb548ecc41 100644 --- a/flash-attention-v100/kernel/paged_kv_utils.cuh +++ b/flash-attention-v100/kernel/paged_kv_utils.cuh @@ -4,78 +4,58 @@ #include #include -__device__ __forceinline__ -size_t resolve_paged_kv_offset( - const int token_idx, - const int page_block_size, - const int* __restrict__ block_table, - const int kv_head_offset, - const int num_kv_heads, - const int head_dim -) { +__device__ __forceinline__ size_t resolve_paged_kv_offset( + const int token_idx, const int page_block_size, + const int* __restrict__ block_table, const int kv_head_offset, + const int num_kv_heads, const int head_dim) { + const int virtual_page_idx = token_idx / page_block_size; + const int page_offset = token_idx % page_block_size; - const int virtual_page_idx = token_idx / page_block_size; - const int page_offset = token_idx % page_block_size; + const int physical_block_idx = block_table[virtual_page_idx]; - const int physical_block_idx = block_table[virtual_page_idx]; + const size_t physical_offset = + (size_t)physical_block_idx * page_block_size * num_kv_heads * head_dim + + (size_t)page_offset * num_kv_heads * head_dim + (size_t)kv_head_offset; - const size_t physical_offset = - (size_t)physical_block_idx * page_block_size * num_kv_heads * head_dim - + (size_t)page_offset * num_kv_heads * head_dim - + (size_t)kv_head_offset; - - return physical_offset; + return physical_offset; } -template -__device__ __forceinline__ -void load_kv_tile_paged( - const __half* __restrict__ kv_cache, - const int* __restrict__ block_table, - const int start_token_idx, - const int num_tokens, - const int kv_head_idx, - const int num_kv_heads, - const int head_dim, - const int page_block_size, - __half* smem_dst, - const int smem_stride, - const int tid, - const int num_threads -) { - constexpr int PER_UINT4 = 8; - const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; - const int smem_stride_uint4 = (smem_stride + PER_UINT4 - 1) / PER_UINT4; - - const int kv_head_offset = kv_head_idx * head_dim; +template +__device__ __forceinline__ void load_kv_tile_paged( + const __half* __restrict__ kv_cache, const int* __restrict__ block_table, + const int start_token_idx, const int num_tokens, const int kv_head_idx, + const int num_kv_heads, const int head_dim, const int page_block_size, + __half* smem_dst, const int smem_stride, const int tid, + const int num_threads) { + constexpr int PER_UINT4 = 8; + const int d_stride_uint4 = (D + PER_UINT4 - 1) / PER_UINT4; + const int smem_stride_uint4 = (smem_stride + PER_UINT4 - 1) / PER_UINT4; - uint4* smem_vec = reinterpret_cast(smem_dst); + const int kv_head_offset = kv_head_idx * head_dim; - #pragma unroll 2 - for (int idx = tid; idx < (num_tokens * d_stride_uint4); idx += num_threads) { - const int token_offset = idx / d_stride_uint4; - const int vec_col = idx % d_stride_uint4; + uint4* smem_vec = reinterpret_cast(smem_dst); - uint4 kv_val = make_uint4(0, 0, 0, 0); +#pragma unroll 2 + for (int idx = tid; idx < (num_tokens * d_stride_uint4); idx += num_threads) { + const int token_offset = idx / d_stride_uint4; + const int vec_col = idx % d_stride_uint4; - if (token_offset < num_tokens && vec_col < d_stride_uint4) { - const int global_token_idx = start_token_idx + token_offset; + uint4 kv_val = make_uint4(0, 0, 0, 0); - const size_t base_offset = resolve_paged_kv_offset( - global_token_idx, - page_block_size, - block_table, - kv_head_offset, - num_kv_heads, - head_dim - ); + if (token_offset < num_tokens && vec_col < d_stride_uint4) { + const int global_token_idx = start_token_idx + token_offset; - const uint4* kv_vec = reinterpret_cast(kv_cache + base_offset); - kv_val = __ldg(&kv_vec[vec_col]); - } + const size_t base_offset = resolve_paged_kv_offset( + global_token_idx, page_block_size, block_table, kv_head_offset, + num_kv_heads, head_dim); - smem_vec[token_offset * smem_stride_uint4 + vec_col] = kv_val; + const uint4* kv_vec = + reinterpret_cast(kv_cache + base_offset); + kv_val = __ldg(&kv_vec[vec_col]); } + + smem_vec[token_offset * smem_stride_uint4 + vec_col] = kv_val; + } } #endif diff --git a/flash-attention-v100/kernel/paged_to_contiguous.cu b/flash-attention-v100/kernel/paged_to_contiguous.cu index b5f321a28a..77a94b13e7 100644 --- a/flash-attention-v100/kernel/paged_to_contiguous.cu +++ b/flash-attention-v100/kernel/paged_to_contiguous.cu @@ -4,215 +4,178 @@ #include __global__ void paged_to_contiguous_stride_aware_kernel( - const __half* __restrict__ paged_cache, - const int* __restrict__ block_table, - const int* __restrict__ seq_lens, - __half* __restrict__ contiguous, - const int64_t block_stride, - const int64_t page_stride, - const int64_t head_stride, - const int block_size, - const int num_heads, - const int head_dim, - const int max_num_blocks -) { - - const int batch_idx = blockIdx.x; - const int seq_len = seq_lens[batch_idx]; - if (seq_len == 0) return; - - int output_offset = 0; - for (int i = 0; i < batch_idx; ++i) { - output_offset += seq_lens[i]; - } - - const int* block_table_seq = block_table + batch_idx * max_num_blocks; - - const int tid = threadIdx.x; - const int num_threads = blockDim.x; - const int elements_per_token = num_heads * head_dim; - - for (int token_idx = 0; token_idx < seq_len; ++token_idx) { - - const int virtual_block_idx = token_idx / block_size; - const int block_offset = token_idx % block_size; - const int physical_block_idx = block_table_seq[virtual_block_idx]; - - const __half* token_base = paged_cache - + physical_block_idx * block_stride - + block_offset * page_stride; - - __half* output_token = contiguous + (output_offset + token_idx) * elements_per_token; - - for (int head_idx = 0; head_idx < num_heads; ++head_idx) { - const __half* head_src = token_base + head_idx * head_stride; - __half* head_dst = output_token + head_idx * head_dim; - - for (int elem_idx = tid; elem_idx < head_dim; elem_idx += num_threads) { - head_dst[elem_idx] = head_src[elem_idx]; - } - } + const __half* __restrict__ paged_cache, const int* __restrict__ block_table, + const int* __restrict__ seq_lens, __half* __restrict__ contiguous, + const int64_t block_stride, const int64_t page_stride, + const int64_t head_stride, const int block_size, const int num_heads, + const int head_dim, const int max_num_blocks) { + const int batch_idx = blockIdx.x; + const int seq_len = seq_lens[batch_idx]; + if (seq_len == 0) return; + + int output_offset = 0; + for (int i = 0; i < batch_idx; ++i) { + output_offset += seq_lens[i]; + } + + const int* block_table_seq = block_table + batch_idx * max_num_blocks; + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + const int elements_per_token = num_heads * head_dim; + + for (int token_idx = 0; token_idx < seq_len; ++token_idx) { + const int virtual_block_idx = token_idx / block_size; + const int block_offset = token_idx % block_size; + const int physical_block_idx = block_table_seq[virtual_block_idx]; + + const __half* token_base = paged_cache + physical_block_idx * block_stride + + block_offset * page_stride; + + __half* output_token = + contiguous + (output_offset + token_idx) * elements_per_token; + + for (int head_idx = 0; head_idx < num_heads; ++head_idx) { + const __half* head_src = token_base + head_idx * head_stride; + __half* head_dst = output_token + head_idx * head_dim; + + for (int elem_idx = tid; elem_idx < head_dim; elem_idx += num_threads) { + head_dst[elem_idx] = head_src[elem_idx]; + } } + } } __global__ void paged_kv_to_contiguous_stride_aware_kernel( const __half* __restrict__ paged_key_cache, const __half* __restrict__ paged_value_cache, - const int* __restrict__ block_table, - const int* __restrict__ seq_lens, - __half* __restrict__ contiguous_key, - __half* __restrict__ contiguous_value, - const int64_t key_block_stride, - const int64_t key_page_stride, - const int64_t key_head_stride, - const int64_t value_block_stride, - const int64_t value_page_stride, - const int64_t value_head_stride, - const int block_size, - const int num_heads, - const int head_dim, - const int max_num_blocks -) { - const int batch_idx = blockIdx.x; - const int seq_len = seq_lens[batch_idx]; - if (seq_len == 0) return; - - int output_offset = 0; - for (int i = 0; i < batch_idx; ++i) { - output_offset += seq_lens[i]; - } - - const int* block_table_seq = block_table + batch_idx * max_num_blocks; - const int tid = threadIdx.x; - const int num_threads = blockDim.x; - const int elements_per_token = num_heads * head_dim; - - for (int token_idx = 0; token_idx < seq_len; ++token_idx) { - const int virtual_block_idx = token_idx / block_size; - const int block_offset = token_idx % block_size; - const int physical_block_idx = block_table_seq[virtual_block_idx]; - - const __half* key_token_base = paged_key_cache - + physical_block_idx * key_block_stride - + block_offset * key_page_stride; - const __half* value_token_base = paged_value_cache - + physical_block_idx * value_block_stride - + block_offset * value_page_stride; - - __half* key_output_token = - contiguous_key + (output_offset + token_idx) * elements_per_token; - __half* value_output_token = - contiguous_value + (output_offset + token_idx) * elements_per_token; - - for (int head_idx = 0; head_idx < num_heads; ++head_idx) { - const __half* key_head_src = key_token_base + head_idx * key_head_stride; - const __half* value_head_src = - value_token_base + head_idx * value_head_stride; - __half* key_head_dst = key_output_token + head_idx * head_dim; - __half* value_head_dst = value_output_token + head_idx * head_dim; - - for (int elem_idx = tid; elem_idx < head_dim; elem_idx += num_threads) { - key_head_dst[elem_idx] = key_head_src[elem_idx]; - value_head_dst[elem_idx] = value_head_src[elem_idx]; - } - } + const int* __restrict__ block_table, const int* __restrict__ seq_lens, + __half* __restrict__ contiguous_key, __half* __restrict__ contiguous_value, + const int64_t key_block_stride, const int64_t key_page_stride, + const int64_t key_head_stride, const int64_t value_block_stride, + const int64_t value_page_stride, const int64_t value_head_stride, + const int block_size, const int num_heads, const int head_dim, + const int max_num_blocks) { + const int batch_idx = blockIdx.x; + const int seq_len = seq_lens[batch_idx]; + if (seq_len == 0) return; + + int output_offset = 0; + for (int i = 0; i < batch_idx; ++i) { + output_offset += seq_lens[i]; + } + + const int* block_table_seq = block_table + batch_idx * max_num_blocks; + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + const int elements_per_token = num_heads * head_dim; + + for (int token_idx = 0; token_idx < seq_len; ++token_idx) { + const int virtual_block_idx = token_idx / block_size; + const int block_offset = token_idx % block_size; + const int physical_block_idx = block_table_seq[virtual_block_idx]; + + const __half* key_token_base = paged_key_cache + + physical_block_idx * key_block_stride + + block_offset * key_page_stride; + const __half* value_token_base = paged_value_cache + + physical_block_idx * value_block_stride + + block_offset * value_page_stride; + + __half* key_output_token = + contiguous_key + (output_offset + token_idx) * elements_per_token; + __half* value_output_token = + contiguous_value + (output_offset + token_idx) * elements_per_token; + + for (int head_idx = 0; head_idx < num_heads; ++head_idx) { + const __half* key_head_src = key_token_base + head_idx * key_head_stride; + const __half* value_head_src = + value_token_base + head_idx * value_head_stride; + __half* key_head_dst = key_output_token + head_idx * head_dim; + __half* value_head_dst = value_output_token + head_idx * head_dim; + + for (int elem_idx = tid; elem_idx < head_dim; elem_idx += num_threads) { + key_head_dst[elem_idx] = key_head_src[elem_idx]; + value_head_dst[elem_idx] = value_head_src[elem_idx]; + } } + } } -torch::Tensor paged_to_contiguous( - torch::Tensor paged_cache, - torch::Tensor block_table, - torch::Tensor seq_lens -) { - const int batch_size = block_table.size(0); - const int max_num_blocks = block_table.size(1); - const int num_blocks = paged_cache.size(0); - const int block_size = paged_cache.size(1); - const int num_heads = paged_cache.size(2); - const int head_dim = paged_cache.size(3); - - const int64_t block_stride = paged_cache.stride(0); - const int64_t page_stride = paged_cache.stride(1); - const int64_t head_stride = paged_cache.stride(2); - - const int total_tokens = batch_size * max_num_blocks * block_size; - - auto contiguous = torch::zeros( - {total_tokens, num_heads, head_dim}, - torch::TensorOptions().dtype(paged_cache.dtype()).device(paged_cache.device()) - ); - - const int threads = 256; - paged_to_contiguous_stride_aware_kernel<<>>( - reinterpret_cast(paged_cache.data_ptr()), - block_table.data_ptr(), - seq_lens.data_ptr(), - reinterpret_cast<__half*>(contiguous.data_ptr()), - block_stride, - page_stride, - head_stride, - block_size, - num_heads, - head_dim, - max_num_blocks - ); - - return contiguous; +torch::Tensor paged_to_contiguous(torch::Tensor paged_cache, + torch::Tensor block_table, + torch::Tensor seq_lens) { + const int batch_size = block_table.size(0); + const int max_num_blocks = block_table.size(1); + const int num_blocks = paged_cache.size(0); + const int block_size = paged_cache.size(1); + const int num_heads = paged_cache.size(2); + const int head_dim = paged_cache.size(3); + + const int64_t block_stride = paged_cache.stride(0); + const int64_t page_stride = paged_cache.stride(1); + const int64_t head_stride = paged_cache.stride(2); + + const int total_tokens = batch_size * max_num_blocks * block_size; + + auto contiguous = torch::zeros({total_tokens, num_heads, head_dim}, + torch::TensorOptions() + .dtype(paged_cache.dtype()) + .device(paged_cache.device())); + + const int threads = 256; + paged_to_contiguous_stride_aware_kernel<<>>( + reinterpret_cast(paged_cache.data_ptr()), + block_table.data_ptr(), seq_lens.data_ptr(), + reinterpret_cast<__half*>(contiguous.data_ptr()), block_stride, + page_stride, head_stride, block_size, num_heads, head_dim, + max_num_blocks); + + return contiguous; } std::vector paged_kv_to_contiguous( - torch::Tensor paged_key_cache, - torch::Tensor paged_value_cache, - torch::Tensor block_table, - torch::Tensor seq_lens -) { - const int batch_size = block_table.size(0); - const int max_num_blocks = block_table.size(1); - const int block_size = paged_key_cache.size(1); - const int num_heads = paged_key_cache.size(2); - const int head_dim = paged_key_cache.size(3); - - const int64_t key_block_stride = paged_key_cache.stride(0); - const int64_t key_page_stride = paged_key_cache.stride(1); - const int64_t key_head_stride = paged_key_cache.stride(2); - const int64_t value_block_stride = paged_value_cache.stride(0); - const int64_t value_page_stride = paged_value_cache.stride(1); - const int64_t value_head_stride = paged_value_cache.stride(2); - - const int total_tokens = batch_size * max_num_blocks * block_size; - - auto opts = torch::TensorOptions() - .dtype(paged_key_cache.dtype()) - .device(paged_key_cache.device()); - auto contiguous_key = torch::zeros({total_tokens, num_heads, head_dim}, opts); - auto contiguous_value = torch::zeros({total_tokens, num_heads, head_dim}, opts); - - const int threads = 256; - paged_kv_to_contiguous_stride_aware_kernel<<>>( - reinterpret_cast(paged_key_cache.data_ptr()), - reinterpret_cast(paged_value_cache.data_ptr()), - block_table.data_ptr(), - seq_lens.data_ptr(), - reinterpret_cast<__half*>(contiguous_key.data_ptr()), - reinterpret_cast<__half*>(contiguous_value.data_ptr()), - key_block_stride, - key_page_stride, - key_head_stride, - value_block_stride, - value_page_stride, - value_head_stride, - block_size, - num_heads, - head_dim, - max_num_blocks - ); - - return {contiguous_key, contiguous_value}; + torch::Tensor paged_key_cache, torch::Tensor paged_value_cache, + torch::Tensor block_table, torch::Tensor seq_lens) { + const int batch_size = block_table.size(0); + const int max_num_blocks = block_table.size(1); + const int block_size = paged_key_cache.size(1); + const int num_heads = paged_key_cache.size(2); + const int head_dim = paged_key_cache.size(3); + + const int64_t key_block_stride = paged_key_cache.stride(0); + const int64_t key_page_stride = paged_key_cache.stride(1); + const int64_t key_head_stride = paged_key_cache.stride(2); + const int64_t value_block_stride = paged_value_cache.stride(0); + const int64_t value_page_stride = paged_value_cache.stride(1); + const int64_t value_head_stride = paged_value_cache.stride(2); + + const int total_tokens = batch_size * max_num_blocks * block_size; + + auto opts = torch::TensorOptions() + .dtype(paged_key_cache.dtype()) + .device(paged_key_cache.device()); + auto contiguous_key = torch::zeros({total_tokens, num_heads, head_dim}, opts); + auto contiguous_value = + torch::zeros({total_tokens, num_heads, head_dim}, opts); + + const int threads = 256; + paged_kv_to_contiguous_stride_aware_kernel<<>>( + reinterpret_cast(paged_key_cache.data_ptr()), + reinterpret_cast(paged_value_cache.data_ptr()), + block_table.data_ptr(), seq_lens.data_ptr(), + reinterpret_cast<__half*>(contiguous_key.data_ptr()), + reinterpret_cast<__half*>(contiguous_value.data_ptr()), + key_block_stride, key_page_stride, key_head_stride, value_block_stride, + value_page_stride, value_head_stride, block_size, num_heads, head_dim, + max_num_blocks); + + return {contiguous_key, contiguous_value}; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("paged_to_contiguous", &paged_to_contiguous, "Paged KV Cache to Contiguous (Stride-Aware)"); - m.def("paged_kv_to_contiguous", - &paged_kv_to_contiguous, - "Paged KV Cache to Contiguous K/V Pair (Stride-Aware)"); + m.def("paged_to_contiguous", &paged_to_contiguous, + "Paged KV Cache to Contiguous (Stride-Aware)"); + m.def("paged_kv_to_contiguous", &paged_kv_to_contiguous, + "Paged KV Cache to Contiguous K/V Pair (Stride-Aware)"); } diff --git a/flash-attention-v100/kernel/paged_to_contiguous_fixed.cu b/flash-attention-v100/kernel/paged_to_contiguous_fixed.cu index c4901cb6d7..aeda1b57cb 100644 --- a/flash-attention-v100/kernel/paged_to_contiguous_fixed.cu +++ b/flash-attention-v100/kernel/paged_to_contiguous_fixed.cu @@ -4,102 +4,85 @@ #include __global__ void paged_to_contiguous_stride_aware_kernel( - const __half* __restrict__ paged_cache, - const int* __restrict__ block_table, - const int* __restrict__ seq_lens, - __half* __restrict__ contiguous, - const int64_t block_stride, - const int64_t page_stride, - const int64_t head_stride, - const int block_size, - const int num_heads, - const int head_dim, - const int max_num_blocks -) { - - const int batch_idx = blockIdx.x; - const int seq_len = seq_lens[batch_idx]; - if (seq_len == 0) return; - - int output_offset = 0; - for (int i = 0; i < batch_idx; ++i) { - output_offset += seq_lens[i]; - } - - const int* block_table_seq = block_table + batch_idx * max_num_blocks; - - const int tid = threadIdx.x; - const int num_threads = blockDim.x; - const int elements_per_token = num_heads * head_dim; - - for (int token_idx = 0; token_idx < seq_len; ++token_idx) { - - const int virtual_block_idx = token_idx / block_size; - const int block_offset = token_idx % block_size; - const int physical_block_idx = block_table_seq[virtual_block_idx]; - - const __half* token_base = paged_cache - + physical_block_idx * block_stride - + block_offset * page_stride; - - __half* output_token = contiguous + (output_offset + token_idx) * elements_per_token; - - for (int head_idx = 0; head_idx < num_heads; ++head_idx) { - const __half* head_src = token_base + head_idx * head_stride; - __half* head_dst = output_token + head_idx * head_dim; - - for (int elem_idx = tid; elem_idx < head_dim; elem_idx += num_threads) { - head_dst[elem_idx] = head_src[elem_idx]; - } - } + const __half* __restrict__ paged_cache, const int* __restrict__ block_table, + const int* __restrict__ seq_lens, __half* __restrict__ contiguous, + const int64_t block_stride, const int64_t page_stride, + const int64_t head_stride, const int block_size, const int num_heads, + const int head_dim, const int max_num_blocks) { + const int batch_idx = blockIdx.x; + const int seq_len = seq_lens[batch_idx]; + if (seq_len == 0) return; + + int output_offset = 0; + for (int i = 0; i < batch_idx; ++i) { + output_offset += seq_lens[i]; + } + + const int* block_table_seq = block_table + batch_idx * max_num_blocks; + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + const int elements_per_token = num_heads * head_dim; + + for (int token_idx = 0; token_idx < seq_len; ++token_idx) { + const int virtual_block_idx = token_idx / block_size; + const int block_offset = token_idx % block_size; + const int physical_block_idx = block_table_seq[virtual_block_idx]; + + const __half* token_base = paged_cache + physical_block_idx * block_stride + + block_offset * page_stride; + + __half* output_token = + contiguous + (output_offset + token_idx) * elements_per_token; + + for (int head_idx = 0; head_idx < num_heads; ++head_idx) { + const __half* head_src = token_base + head_idx * head_stride; + __half* head_dst = output_token + head_idx * head_dim; + + for (int elem_idx = tid; elem_idx < head_dim; elem_idx += num_threads) { + head_dst[elem_idx] = head_src[elem_idx]; + } } + } } -torch::Tensor paged_to_contiguous( - torch::Tensor paged_cache, - torch::Tensor block_table, - torch::Tensor seq_lens -) { - const int batch_size = block_table.size(0); - const int max_num_blocks = block_table.size(1); - const int num_blocks = paged_cache.size(0); - const int block_size = paged_cache.size(1); - const int num_heads = paged_cache.size(2); - const int head_dim = paged_cache.size(3); - - const int64_t block_stride = paged_cache.stride(0); - const int64_t page_stride = paged_cache.stride(1); - const int64_t head_stride = paged_cache.stride(2); - - int total_tokens = 0; - auto seq_lens_cpu = seq_lens.cpu(); - for (int i = 0; i < batch_size; ++i) { - total_tokens += seq_lens_cpu[i].item(); - } - - auto contiguous = torch::zeros( - {total_tokens, num_heads, head_dim}, - torch::TensorOptions().dtype(paged_cache.dtype()).device(paged_cache.device()) - ); - - const int threads = 256; - paged_to_contiguous_stride_aware_kernel<<>>( - reinterpret_cast(paged_cache.data_ptr()), - block_table.data_ptr(), - seq_lens.data_ptr(), - reinterpret_cast<__half*>(contiguous.data_ptr()), - block_stride, - page_stride, - head_stride, - block_size, - num_heads, - head_dim, - max_num_blocks - ); - - return contiguous; +torch::Tensor paged_to_contiguous(torch::Tensor paged_cache, + torch::Tensor block_table, + torch::Tensor seq_lens) { + const int batch_size = block_table.size(0); + const int max_num_blocks = block_table.size(1); + const int num_blocks = paged_cache.size(0); + const int block_size = paged_cache.size(1); + const int num_heads = paged_cache.size(2); + const int head_dim = paged_cache.size(3); + + const int64_t block_stride = paged_cache.stride(0); + const int64_t page_stride = paged_cache.stride(1); + const int64_t head_stride = paged_cache.stride(2); + + int total_tokens = 0; + auto seq_lens_cpu = seq_lens.cpu(); + for (int i = 0; i < batch_size; ++i) { + total_tokens += seq_lens_cpu[i].item(); + } + + auto contiguous = torch::zeros({total_tokens, num_heads, head_dim}, + torch::TensorOptions() + .dtype(paged_cache.dtype()) + .device(paged_cache.device())); + + const int threads = 256; + paged_to_contiguous_stride_aware_kernel<<>>( + reinterpret_cast(paged_cache.data_ptr()), + block_table.data_ptr(), seq_lens.data_ptr(), + reinterpret_cast<__half*>(contiguous.data_ptr()), block_stride, + page_stride, head_stride, block_size, num_heads, head_dim, + max_num_blocks); + + return contiguous; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("paged_to_contiguous", &paged_to_contiguous, "Paged KV Cache to Contiguous (Stride-Aware)"); + m.def("paged_to_contiguous", &paged_to_contiguous, + "Paged KV Cache to Contiguous (Stride-Aware)"); } diff --git a/flash-attention-v100/kernel/paged_to_contiguous_old.cu b/flash-attention-v100/kernel/paged_to_contiguous_old.cu index 0017127d71..9845ab19fb 100644 --- a/flash-attention-v100/kernel/paged_to_contiguous_old.cu +++ b/flash-attention-v100/kernel/paged_to_contiguous_old.cu @@ -4,85 +4,75 @@ #include __global__ void paged_to_contiguous_kernel( - const __half* __restrict__ paged_cache, - const int* __restrict__ block_table, - const int* __restrict__ seq_lens, - __half* __restrict__ contiguous, - const int block_size, - const int num_heads, - const int head_dim, - const int max_num_blocks -) { - - const int batch_idx = blockIdx.x; - const int seq_len = seq_lens[batch_idx]; - if (seq_len == 0) return; - - int output_offset = 0; - for (int i = 0; i < batch_idx; ++i) { - output_offset += seq_lens[i]; - } - - const int* block_table_seq = block_table + batch_idx * max_num_blocks; - - const int tid = threadIdx.x; - const int num_threads = blockDim.x; - const int elements_per_token = num_heads * head_dim; - - for (int token_idx = 0; token_idx < seq_len; ++token_idx) { - - const int virtual_block_idx = token_idx / block_size; - const int block_offset = token_idx % block_size; - const int physical_block_idx = block_table_seq[virtual_block_idx]; - - const size_t paged_offset = (size_t)physical_block_idx * block_size * elements_per_token - + (size_t)block_offset * elements_per_token; - - const size_t cont_offset = (size_t)(output_offset + token_idx) * elements_per_token; - - for (int elem_idx = tid; elem_idx < elements_per_token; elem_idx += num_threads) { - contiguous[cont_offset + elem_idx] = paged_cache[paged_offset + elem_idx]; - } + const __half* __restrict__ paged_cache, const int* __restrict__ block_table, + const int* __restrict__ seq_lens, __half* __restrict__ contiguous, + const int block_size, const int num_heads, const int head_dim, + const int max_num_blocks) { + const int batch_idx = blockIdx.x; + const int seq_len = seq_lens[batch_idx]; + if (seq_len == 0) return; + + int output_offset = 0; + for (int i = 0; i < batch_idx; ++i) { + output_offset += seq_lens[i]; + } + + const int* block_table_seq = block_table + batch_idx * max_num_blocks; + + const int tid = threadIdx.x; + const int num_threads = blockDim.x; + const int elements_per_token = num_heads * head_dim; + + for (int token_idx = 0; token_idx < seq_len; ++token_idx) { + const int virtual_block_idx = token_idx / block_size; + const int block_offset = token_idx % block_size; + const int physical_block_idx = block_table_seq[virtual_block_idx]; + + const size_t paged_offset = + (size_t)physical_block_idx * block_size * elements_per_token + + (size_t)block_offset * elements_per_token; + + const size_t cont_offset = + (size_t)(output_offset + token_idx) * elements_per_token; + + for (int elem_idx = tid; elem_idx < elements_per_token; + elem_idx += num_threads) { + contiguous[cont_offset + elem_idx] = paged_cache[paged_offset + elem_idx]; } + } } -torch::Tensor paged_to_contiguous( - torch::Tensor paged_cache, - torch::Tensor block_table, - torch::Tensor seq_lens -) { - const int batch_size = block_table.size(0); - const int max_num_blocks = block_table.size(1); - const int block_size = paged_cache.size(1); - const int num_heads = paged_cache.size(2); - const int head_dim = paged_cache.size(3); - - int total_tokens = 0; - auto seq_lens_cpu = seq_lens.cpu(); - for (int i = 0; i < batch_size; ++i) { - total_tokens += seq_lens_cpu[i].item(); - } - - auto contiguous = torch::zeros( - {total_tokens, num_heads, head_dim}, - torch::TensorOptions().dtype(paged_cache.dtype()).device(paged_cache.device()) - ); - - const int threads = 256; - paged_to_contiguous_kernel<<>>( - reinterpret_cast(paged_cache.data_ptr()), - block_table.data_ptr(), - seq_lens.data_ptr(), - reinterpret_cast<__half*>(contiguous.data_ptr()), - block_size, - num_heads, - head_dim, - max_num_blocks - ); - - return contiguous; +torch::Tensor paged_to_contiguous(torch::Tensor paged_cache, + torch::Tensor block_table, + torch::Tensor seq_lens) { + const int batch_size = block_table.size(0); + const int max_num_blocks = block_table.size(1); + const int block_size = paged_cache.size(1); + const int num_heads = paged_cache.size(2); + const int head_dim = paged_cache.size(3); + + int total_tokens = 0; + auto seq_lens_cpu = seq_lens.cpu(); + for (int i = 0; i < batch_size; ++i) { + total_tokens += seq_lens_cpu[i].item(); + } + + auto contiguous = torch::zeros({total_tokens, num_heads, head_dim}, + torch::TensorOptions() + .dtype(paged_cache.dtype()) + .device(paged_cache.device())); + + const int threads = 256; + paged_to_contiguous_kernel<<>>( + reinterpret_cast(paged_cache.data_ptr()), + block_table.data_ptr(), seq_lens.data_ptr(), + reinterpret_cast<__half*>(contiguous.data_ptr()), block_size, + num_heads, head_dim, max_num_blocks); + + return contiguous; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("paged_to_contiguous", &paged_to_contiguous, "Paged KV Cache to Contiguous"); + m.def("paged_to_contiguous", &paged_to_contiguous, + "Paged KV Cache to Contiguous"); } diff --git a/flash_qla/ops/gated_delta_rule/chunk/sm70/fused_fwd.py b/flash_qla/ops/gated_delta_rule/chunk/sm70/fused_fwd.py index 3a16c15e1e..6ca8bfc59e 100644 --- a/flash_qla/ops/gated_delta_rule/chunk/sm70/fused_fwd.py +++ b/flash_qla/ops/gated_delta_rule/chunk/sm70/fused_fwd.py @@ -4,6 +4,7 @@ from __future__ import annotations import importlib +import importlib.util import os from pathlib import Path @@ -15,6 +16,7 @@ _PREBUILT_EXTENSION = ( "flash_qla.ops.gated_delta_rule.chunk.sm70.flash_qla_sm70_gdn_strided" ) +_PREBUILT_EXTENSION_PATH_ENV = "FLASH_QLA_SM70_PREBUILT_EXTENSION_PATH" # Avoid CUDA-device auto-detection while distributed workers are initializing. # The fallback JIT supports both architectures implemented by this source file. @@ -36,6 +38,25 @@ def _load_ext(): if not torch.cuda.is_available(): raise RuntimeError("SM70 FlashQLA backend requires CUDA.") + prebuilt_path = os.environ.get(_PREBUILT_EXTENSION_PATH_ENV) + if prebuilt_path: + path = Path(prebuilt_path).expanduser() + try: + if not path.is_file(): + raise FileNotFoundError(path) + spec = importlib.util.spec_from_file_location(_PREBUILT_EXTENSION, path) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load extension spec from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + _EXT = module + return _EXT + except Exception as exc: + _EXT_LOAD_ERROR = exc + raise RuntimeError( + f"Failed to load {_PREBUILT_EXTENSION_PATH_ENV}={path}." + ) from exc + try: _EXT = importlib.import_module(_PREBUILT_EXTENSION) return _EXT diff --git a/requirements/common.txt b/requirements/common.txt index 0384ca2f96..27ee2a5918 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -3,31 +3,31 @@ cachetools psutil sentencepiece # Required for LLaMA tokenizer. numpy -requests >= 2.26.0 +requests >= 2.34.2 tqdm blake3 py-cpuinfo -transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0 +transformers >= 5.15.1 tokenizers >= 0.21.1 # Required for fast incremental detokenization. -safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 +safetensors >= 0.8.0 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994 -fastapi[standard] >= 0.115.0, < 0.136.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint. -aiohttp >= 3.13.3 +fastapi[standard] >= 0.141.1, < 0.142.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint. +aiohttp >= 3.14.3 openai >= 2.0.0 # For Responses API with reasoning content -pydantic >= 2.12.0, < 2.13.0 -prometheus_client >= 0.18.0 +pydantic >= 2.13.4, < 2.14.0 +prometheus_client >= 0.26.0 pillow # Required for image processing prometheus-fastapi-instrumentator >= 7.0.0, < 8.0.0 -tiktoken >= 0.6.0 # Required for DBRX tokenizer +tiktoken >= 0.14.0 # Required for DBRX tokenizer lm-format-enforcer == 0.11.3 llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" outlines_core == 0.2.14 # required for outlines backend disk cache diskcache == 5.6.3 -lark == 1.2.2 +lark == 1.3.1 xgrammar >= 0.2.0, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le" -typing_extensions >= 4.10 -filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317 +typing_extensions >= 4.16.0 +filelock >= 3.32.3 # need to contain https://github.com/tox-dev/filelock/pull/317 partial-json-parser # used for parsing partial JSON outputs pyzmq >= 25.0.0 msgspec @@ -52,7 +52,7 @@ openai-harmony >= 0.0.3 # Required for gpt-oss anthropic >= 0.71.0 model-hosting-container-standards >= 0.1.14, < 1.0.0 mcp -opentelemetry-sdk >= 1.27.0 -opentelemetry-api >= 1.27.0 +opentelemetry-sdk >= 1.44.0 +opentelemetry-api >= 1.44.0 opentelemetry-exporter-otlp >= 1.27.0 opentelemetry-semantic-conventions-ai >= 0.4.1 diff --git a/requirements/cpu.txt b/requirements/cpu.txt index 5ec338af73..e22cde62d2 100644 --- a/requirements/cpu.txt +++ b/requirements/cpu.txt @@ -4,7 +4,7 @@ setuptools==77.0.3 # this version can reuse CMake build dir -numba == 0.65.0; platform_machine != "s390x" # Required for N-gram speculative decoding +numba == 0.67.0; platform_machine != "s390x" # Required for N-gram speculative decoding # Dependencies for CPUs torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 8d92a144c1..52161d05f1 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -1,7 +1,7 @@ # Common dependencies -r common.txt -numba == 0.65.0 # Required for N-gram speculative decoding +numba == 0.67.0 # Required for N-gram speculative decoding # Dependencies for NVIDIA GPUs torch==2.10.0 @@ -18,14 +18,14 @@ tilelang==0.1.10 nvidia-cudnn-frontend>=1.13.0,<1.19.0 # Required for faster safetensors model loading -fastsafetensors >= 0.2.2 +fastsafetensors >= 0.3.3 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) -nvidia-cutlass-dsl[cu13]==4.5.2 +nvidia-cutlass-dsl[cu13]==4.7.0 quack-kernels>=0.3.3 # Tokenspeed_MLA for faster mla with spec decode -tokenspeed-mla==0.1.2 +tokenspeed-mla==0.2.5 # Humming kernels for quantization gemm humming-kernels[cu13]==0.1.2 diff --git a/requirements/docs.txt b/requirements/docs.txt index a7affa70ea..ffe58c6366 100644 --- a/requirements/docs.txt +++ b/requirements/docs.txt @@ -30,7 +30,7 @@ gitdb==4.0.12 # via gitpython gitpython==3.1.50 # via mkdocs-git-revision-date-localized-plugin -griffelib==2.0.2 +griffelib==2.2.0 # via mkdocstrings-python htmlmin2==0.1.13 # via mkdocs-minify-plugin @@ -109,7 +109,7 @@ msgspec==0.21.1 # via -r requirements/docs.in natsort==8.4.0 # via mkdocs-awesome-nav -packaging==26.2 +packaging==26.3 # via # mkdocs # properdocs @@ -157,7 +157,7 @@ regex==2026.4.4 # via -r requirements/docs.in requests==2.33.1 # via mkdocs-material -ruff==0.15.12 +ruff==0.16.4 # via -r requirements/docs.in selectolax==0.4.8 # via mkdocs-glightbox @@ -167,7 +167,6 @@ smmap==5.0.3 # via gitdb typing-extensions==4.15.0 # via - # mkdocstrings-python # pydantic # pydantic-core # typing-inspection diff --git a/requirements/rocm.txt b/requirements/rocm.txt index 0520f4ca1e..b99e5820ea 100644 --- a/requirements/rocm.txt +++ b/requirements/rocm.txt @@ -2,22 +2,22 @@ -r common.txt # The version of gRPC libraries should be consistent with each other -grpcio==1.78.0 -grpcio-reflection==1.78.0 +grpcio==1.83.0 +grpcio-reflection==1.83.0 -numba == 0.65.0 # Required for N-gram speculative decoding +numba == 0.67.0 # Required for N-gram speculative decoding # Dependencies for AMD GPUs datasets peft pytest-asyncio -tensorizer==2.10.1 +tensorizer==2.12.1 packaging>=24.2 setuptools>=77.0.3,<80.0.0 setuptools-scm>=8 setuptools-rust>=1.9.0 -runai-model-streamer[s3,gcs,azure]==0.15.7 -conch-triton-kernels==1.2.1 +runai-model-streamer[s3,gcs,azure]==0.16.1 +conch-triton-kernels==1.3 timm>=1.0.17 # amd-quark: required for Quark quantization on ROCm # To be consistent with test_quark.py diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index e0981c6fdc..f9e2114a48 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -1,6 +1,6 @@ # testing pytest -tensorizer==2.10.1 +tensorizer==2.12.1 pytest-forked pytest-asyncio pytest-rerunfailures @@ -18,7 +18,7 @@ httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test -peft>=0.18.1 # required for phi-4-mm test +peft>=0.20.0 # required for phi-4-mm test pqdm ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests sentence-transformers>=5.2.0 # required for embedding tests @@ -38,11 +38,11 @@ opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.5.3 +transformers==5.15.1 tokenizers==0.22.2 schemathesis>=3.39.15 # Required for openai schema test. # quantization -bitsandbytes==0.49.2 +bitsandbytes==0.50.1 buildkite-test-collector==0.1.9 @@ -50,16 +50,16 @@ genai_perf>=0.0.8 tritonclient>=2.51.0 # The version of gRPC libraries should be consistent with each other -grpcio==1.78.0 -grpcio-reflection==1.78.0 +grpcio==1.83.0 +grpcio-reflection==1.83.0 -arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test -numba == 0.65.0 # Required for N-gram speculative decoding +arctic-inference == 0.2.0; platform_machine == "x86_64" # Required for suffix decoding test +numba == 0.67.0 # Required for N-gram speculative decoding numpy -runai-model-streamer[s3,gcs,azure]==0.15.7 -fastsafetensors>=0.2.2; platform_machine == "x86_64" # 0.2.2 contains important fixes for multi-GPU mem usage +runai-model-streamer[s3,gcs,azure]==0.16.1 +fastsafetensors>=0.3.3; platform_machine == "x86_64" instanttensor>=0.1.5; platform_machine == "x86_64" -pydantic>=2.12 # 2.11 leads to error on python 3.13 +pydantic>=2.13.4 # 2.11 leads to error on python 3.13 decord==0.6.0; platform_machine == "x86_64" # terratorch is temporarily disabled while PyPI has the `lightning` package # in `quarantined` status (every published terratorch version transitively diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 7470b91afd..184c180a5e 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -6,7 +6,7 @@ accelerate==1.13.0 # via peft aiohappyeyeballs==2.6.1 # via aiohttp -aiohttp==3.13.3 +aiohttp==3.14.3 # via # -c requirements/common.txt # aiohttp-cors @@ -31,7 +31,7 @@ anyio==4.6.2.post1 # via # httpx # starlette -arctic-inference==0.1.1 +arctic-inference==0.2.0 # via -r requirements/test/cuda.in argcomplete==3.5.1 # via datamodel-code-generator @@ -60,9 +60,9 @@ backoff==2.2.1 # via # -r requirements/test/cuda.in # schemathesis -bitsandbytes==0.49.2 +bitsandbytes==0.50.1 # via -r requirements/test/cuda.in -black==24.10.0 +black==26.5.1 # via datamodel-code-generator blobfile==3.0.0 # via -r requirements/test/cuda.in @@ -91,7 +91,7 @@ cffi==2.0.0 # via # cryptography # soundfile -chardet==5.2.0 +chardet==6.0.0.post1 # via mbstrdecoder charset-normalizer==3.4.0 # via requests @@ -181,7 +181,7 @@ et-xmlfile==2.0.0 # via openpyxl evaluate==0.4.3 # via lm-eval -fastapi==0.128.0 +fastapi==0.141.1 # via # -c requirements/common.txt # gpt-oss @@ -189,11 +189,11 @@ fastparquet==2024.11.0 # via genai-perf fastrlock==0.8.2 # via cupy-cuda12x -fastsafetensors==0.2.2 +fastsafetensors==0.3.3 # via # -c requirements/cuda.txt # -r requirements/test/cuda.in -filelock==3.16.1 +filelock==3.32.4 # via # -c requirements/common.txt # blobfile @@ -225,7 +225,7 @@ genai-perf==0.0.16 # via -r requirements/test/cuda.in genson==1.3.0 # via datamodel-code-generator -google-api-core==2.24.2 +google-api-core==2.34.0 # via # google-cloud-core # google-cloud-storage @@ -246,7 +246,7 @@ google-crc32c==1.7.1 # google-resumable-media google-resumable-media==2.7.2 # via google-cloud-storage -googleapis-common-protos==1.70.0 +googleapis-common-protos==1.75.2 # via google-api-core gpt-oss==0.0.8 # via -r requirements/test/cuda.in @@ -254,12 +254,12 @@ graphql-core==3.2.6 # via hypothesis-graphql greenlet==3.2.3 # via sqlalchemy -grpcio==1.78.0 +grpcio==1.83.0 # via # -r requirements/test/cuda.in # grpcio-reflection # ray -grpcio-reflection==1.78.0 +grpcio-reflection==1.83.0 # via -r requirements/test/cuda.in h11==0.14.0 # via @@ -323,8 +323,6 @@ imagehash==4.3.2 # via -r requirements/test/cuda.in imageio==2.37.0 # via scikit-image -importlib-metadata==8.7.0 - # via opentelemetry-api inflect==5.6.2 # via datamodel-code-generator iniconfig==2.0.0 @@ -380,7 +378,7 @@ libnacl==2.1.0 # via tensorizer librosa==0.10.2.post1 # via -r requirements/test/cuda.in -llvmlite==0.47.0 +llvmlite==0.49.0 # via numba lm-eval==0.4.12 # via -r requirements/test/cuda.in @@ -400,7 +398,7 @@ markupsafe==3.0.1 # werkzeug matplotlib==3.9.2 # via -r requirements/test/cuda.in -mbstrdecoder==1.1.3 +mbstrdecoder==1.1.5 # via # dataproperty # pytablewriter @@ -445,7 +443,7 @@ nltk==3.9.1 # via rouge-score num2words==0.5.14 # via -r requirements/test/cuda.in -numba==0.65.0 +numba==0.67.0 # via # -c requirements/cuda.txt # -r requirements/test/cuda.in @@ -488,6 +486,7 @@ numpy==2.2.6 # scikit-learn # scipy # segmentation-models-pytorch + # sentence-transformers # soxr # statsmodels # tensorizer @@ -553,22 +552,22 @@ opencv-python-headless==4.13.0.90 # mistral-common openpyxl==3.1.5 # via -r requirements/test/cuda.in -opentelemetry-api==1.35.0 +opentelemetry-api==1.44.0 # via # -c requirements/common.txt # opentelemetry-exporter-prometheus # opentelemetry-sdk # opentelemetry-semantic-conventions -opentelemetry-exporter-prometheus==0.56b0 +opentelemetry-exporter-prometheus==0.65b0 # via ray -opentelemetry-proto==1.35.0 +opentelemetry-proto==1.44.0 # via ray -opentelemetry-sdk==1.35.0 +opentelemetry-sdk==1.44.0 # via # -c requirements/common.txt # opentelemetry-exporter-prometheus # ray -opentelemetry-semantic-conventions==0.56b0 +opentelemetry-semantic-conventions==0.65b0 # via opentelemetry-sdk optuna==3.6.1 # via genai-perf @@ -597,7 +596,7 @@ packaging==26.3 # statsmodels # transformers # typepy -pandas==2.2.3 +pandas==2.3.3 # via # datasets # evaluate @@ -610,13 +609,13 @@ pathvalidate==3.2.1 # via pytablewriter patsy==1.0.1 # via statsmodels -peft==0.18.1 +peft==0.20.0 # via -r requirements/test/cuda.in perceptron==0.1.4 # via -r requirements/test/cuda.in perf-analyzer==0.1.0 # via genai-perf -pillow==10.4.0 +pillow==12.3.0 # via # genai-perf # imagehash @@ -648,7 +647,7 @@ portalocker==2.10.1 # via sacrebleu pqdm==0.2.0 # via -r requirements/test/cuda.in -prometheus-client==0.22.0 +prometheus-client==0.26.0 # via # -c requirements/common.txt # opentelemetry-exporter-prometheus @@ -657,9 +656,9 @@ propcache==0.2.0 # via # aiohttp # yarl -proto-plus==1.26.1 +proto-plus==1.28.4 # via google-api-core -protobuf==6.33.6 +protobuf==7.36.0 # via # -c requirements/common.txt # google-api-core @@ -694,7 +693,7 @@ pycparser==2.22 # via cffi pycryptodomex==3.22.0 # via blobfile -pydantic==2.12.0 +pydantic==2.13.4 # via # -c requirements/common.txt # -r requirements/test/cuda.in @@ -707,7 +706,7 @@ pydantic==2.12.0 # openai-harmony # pydantic-extra-types # ray -pydantic-core==2.41.1 +pydantic-core==2.46.4 # via pydantic pydantic-extra-types==2.10.5 # via mistral-common @@ -762,6 +761,8 @@ python-dateutil==2.9.0.post0 # typepy python-rapidjson==1.20 # via tritonclient +pytokens==0.4.1 + # via black pytrec-eval-terrier==0.5.7 # via mteb pytz==2024.2 @@ -803,7 +804,7 @@ regex==2026.2.28 # sacrebleu # tiktoken # transformers -requests==2.32.3 +requests==2.34.2 # via # -c requirements/common.txt # azure-core @@ -844,19 +845,19 @@ rpds-py==0.20.1 # referencing rsa==4.9.1 # via google-auth -runai-model-streamer==0.15.7 +runai-model-streamer==0.16.1 # via -r requirements/test/cuda.in -runai-model-streamer-azure==0.15.7 +runai-model-streamer-azure==0.16.1 # via runai-model-streamer -runai-model-streamer-gcs==0.15.7 +runai-model-streamer-gcs==0.16.1 # via runai-model-streamer -runai-model-streamer-s3==0.15.7 +runai-model-streamer-s3==0.16.1 # via runai-model-streamer s3transfer==0.10.3 # via boto3 sacrebleu==2.4.3 # via lm-eval -safetensors==0.7.0 +safetensors==0.8.0 # via # -c requirements/common.txt # accelerate @@ -890,7 +891,7 @@ scipy==1.13.1 # vocos segmentation-models-pytorch==0.5.0 # via -r requirements/test/cuda.in -sentence-transformers==5.2.0 +sentence-transformers==6.0.0 # via # -r requirements/test/cuda.in # mteb @@ -963,7 +964,7 @@ tenacity==9.1.2 # gpt-oss # lm-eval # plotly -tensorizer==2.10.1 +tensorizer==2.12.1 # via -r requirements/test/cuda.in termcolor==3.1.0 # via gpt-oss @@ -971,7 +972,7 @@ threadpoolctl==3.5.0 # via scikit-learn tifffile==2025.3.30 # via scikit-image -tiktoken==0.12.0 +tiktoken==0.14.0 # via # -c requirements/common.txt # gpt-oss @@ -986,6 +987,7 @@ tokenizers==0.22.2 # via # -c requirements/common.txt # -r requirements/test/cuda.in + # sentence-transformers # transformers tomli==2.2.1 # via schemathesis @@ -1039,7 +1041,7 @@ tqdm==4.67.3 # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.5.3 +transformers==5.15.1 # via # -c requirements/common.txt # -r requirements/test/cuda.in @@ -1066,9 +1068,10 @@ typer==0.15.2 # transformers types-python-dateutil==2.9.0.20241206 # via arrow -typing-extensions==4.15.0 +typing-extensions==4.16.0 # via # -c requirements/common.txt + # aiohttp # aiosignal # albumentations # alembic @@ -1097,7 +1100,9 @@ typing-extensions==4.15.0 # typer # typing-inspection typing-inspection==0.4.2 - # via pydantic + # via + # fastapi + # pydantic tzdata==2024.2 # via pandas uri-template==1.3.0 @@ -1120,13 +1125,13 @@ vocos==0.1.0 # via -r requirements/test/cuda.in wcwidth==0.2.13 # via ftfy -webcolors==24.11.1 +webcolors==25.10.0 # via jsonschema werkzeug==3.1.3 # via schemathesis word2number==1.1 # via lm-eval -wrapt==1.17.2 +wrapt==2.3.0 # via smart-open xxhash==3.5.0 # via @@ -1136,5 +1141,3 @@ yarl==1.17.1 # via # aiohttp # schemathesis -zipp==3.23.0 - # via importlib-metadata diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 97e0658fb1..3910081624 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -2,7 +2,7 @@ # testing pytest -tensorizer==2.10.1 +tensorizer==2.12.1 pytest-forked pytest-asyncio pytest-rerunfailures @@ -20,7 +20,7 @@ httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test -peft>=0.15.0 # required for phi-4-mm test +peft>=0.20.0 # required for phi-4-mm test pqdm ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests sentence-transformers>=5.2.0 # required for embedding tests @@ -37,11 +37,11 @@ opencv-python-headless>=4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.5.3 +transformers==5.15.1 tokenizers==0.22.2 schemathesis>=3.39.15 # Required for openai schema test # quantization -bitsandbytes==0.49.2 +bitsandbytes==0.50.1 buildkite-test-collector==0.1.9 tilelang==0.1.10 @@ -49,16 +49,16 @@ genai_perf>=0.0.8 tritonclient>=2.51.0 # The version of gRPC libraries should be consistent with each other -grpcio==1.78.0 -grpcio-reflection==1.78.0 +grpcio==1.83.0 +grpcio-reflection==1.83.0 -arctic-inference==0.1.1 # Required for suffix decoding test -numba==0.65.0 # Required for N-gram speculative decoding +arctic-inference==0.2.0 # Required for suffix decoding test +numba==0.67.0 # Required for N-gram speculative decoding numpy -runai-model-streamer[s3,gcs,azure]==0.15.7 -fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels +runai-model-streamer[s3,gcs,azure]==0.16.1 +fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.3.3 # PyPI only ships CUDA wheels instanttensor>=0.1.5 -pydantic>=2.12 # 2.11 leads to error on python 3.13 +pydantic>=2.13.4 # 2.11 leads to error on python 3.13 decord==0.6.0 # Prithvi tests diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index c39f268709..5a40ff1c4a 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -6,7 +6,7 @@ accelerate==1.13.0 # via peft aiohappyeyeballs==2.6.1 # via aiohttp -aiohttp==3.13.3 +aiohttp==3.14.3 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -46,7 +46,7 @@ apache-tvm-ffi==0.1.10 # via # tilelang # xgrammar -arctic-inference==0.1.1 +arctic-inference==0.2.0 # via -r requirements/test/rocm.in argcomplete==3.6.3 # via datamodel-code-generator @@ -76,7 +76,7 @@ backoff==2.2.1 # via # -r requirements/test/rocm.in # schemathesis -bitsandbytes==0.49.2 +bitsandbytes==0.50.1 # via -r requirements/test/rocm.in black==26.3.1 # via datamodel-code-generator @@ -226,23 +226,25 @@ et-xmlfile==2.0.0 # via openpyxl evaluate==0.4.6 # via lm-eval -fastapi==0.135.2 +fastapi==0.141.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt # gpt-oss # model-hosting-container-standards -fastapi-cli==0.0.24 +fastapi-cli==0.0.32 # via fastapi fastapi-cloud-cli==0.16.1 # via fastapi-cli fastar==0.10.0 - # via fastapi-cloud-cli + # via + # fastapi + # fastapi-cloud-cli fastparquet==2026.3.0 # via genai-perf -fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c +fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@daf6cc3fd803ecbf85865062c860122bc42ffefa # via -r requirements/test/rocm.in -filelock==3.25.2 +filelock==3.32.4 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -280,7 +282,7 @@ gguf==0.18.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -google-api-core==2.30.0 +google-api-core==2.34.0 # via # google-cloud-core # google-cloud-storage @@ -301,7 +303,7 @@ google-crc32c==1.8.0 # google-resumable-media google-resumable-media==2.8.0 # via google-cloud-storage -googleapis-common-protos==1.73.0 +googleapis-common-protos==1.75.2 # via # google-api-core # opentelemetry-exporter-otlp-proto-grpc @@ -312,14 +314,14 @@ graphql-core==3.2.8 # via hypothesis-graphql greenlet==3.3.2 # via sqlalchemy -grpcio==1.78.0 +grpcio==1.83.0 # via # -c requirements/rocm.txt # -r requirements/test/rocm.in # grpcio-reflection # opentelemetry-exporter-otlp-proto-grpc # ray -grpcio-reflection==1.78.0 +grpcio-reflection==1.83.0 # via # -c requirements/rocm.txt # -r requirements/test/rocm.in @@ -397,8 +399,6 @@ imagehash==4.3.2 # via -r requirements/test/rocm.in imageio==2.37.3 # via scikit-image -importlib-metadata==8.7.1 - # via opentelemetry-api inflect==7.5.0 # via datamodel-code-generator iniconfig==2.3.0 @@ -455,7 +455,7 @@ kaleido==1.0.0 # via genai-perf kiwisolver==1.5.0 # via matplotlib -lark==1.2.2 +lark==1.3.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -471,7 +471,7 @@ llguidance==1.7.5 # via # -c requirements/common.txt # -r requirements/test/../common.txt -llvmlite==0.47.0 +llvmlite==0.49.0 # via numba lm-eval==0.4.12 # via -r requirements/test/rocm.in @@ -564,7 +564,7 @@ nltk==3.9.3 # via rouge-score num2words==0.5.14 # via -r requirements/test/rocm.in -numba==0.65.0 +numba==0.67.0 # via # -c requirements/rocm.txt # -r requirements/test/rocm.in @@ -647,7 +647,7 @@ opencv-python-headless==4.13.0.92 # mistral-common openpyxl==3.1.5 # via -r requirements/test/rocm.in -opentelemetry-api==1.40.0 +opentelemetry-api==1.44.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -656,27 +656,27 @@ opentelemetry-api==1.40.0 # opentelemetry-exporter-prometheus # opentelemetry-sdk # opentelemetry-semantic-conventions -opentelemetry-exporter-otlp==1.40.0 +opentelemetry-exporter-otlp==1.44.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt -opentelemetry-exporter-otlp-proto-common==1.40.0 +opentelemetry-exporter-otlp-proto-common==1.44.0 # via # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http -opentelemetry-exporter-otlp-proto-grpc==1.40.0 +opentelemetry-exporter-otlp-proto-grpc==1.44.0 # via opentelemetry-exporter-otlp -opentelemetry-exporter-otlp-proto-http==1.40.0 +opentelemetry-exporter-otlp-proto-http==1.44.0 # via opentelemetry-exporter-otlp -opentelemetry-exporter-prometheus==0.61b0 +opentelemetry-exporter-prometheus==0.65b0 # via ray -opentelemetry-proto==1.40.0 +opentelemetry-proto==1.44.0 # via # opentelemetry-exporter-otlp-proto-common # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http # ray -opentelemetry-sdk==1.40.0 +opentelemetry-sdk==1.44.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -685,7 +685,7 @@ opentelemetry-sdk==1.40.0 # opentelemetry-exporter-prometheus # opentelemetry-semantic-conventions-ai # ray -opentelemetry-semantic-conventions==0.61b0 +opentelemetry-semantic-conventions==0.65b0 # via # opentelemetry-sdk # opentelemetry-semantic-conventions-ai @@ -743,7 +743,7 @@ pathvalidate==3.3.1 # via pytablewriter patsy==1.0.2 # via statsmodels -peft==0.18.1 +peft==0.20.0 # via -r requirements/test/rocm.in perceptron==0.1.4 # via -r requirements/test/rocm.in @@ -785,7 +785,7 @@ portalocker==3.2.0 # via sacrebleu pqdm==0.2.0 # via -r requirements/test/rocm.in -prometheus-client==0.24.1 +prometheus-client==0.26.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -800,9 +800,9 @@ propcache==0.4.1 # via # aiohttp # yarl -proto-plus==1.27.1 +proto-plus==1.28.4 # via google-api-core -protobuf==6.33.6 +protobuf==7.36.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -842,7 +842,7 @@ pycparser==3.0 # via cffi pycryptodomex==3.23.0 # via blobfile -pydantic==2.12.5 +pydantic==2.13.4 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -865,7 +865,7 @@ pydantic==2.12.5 # pydantic-settings # ray # xgrammar -pydantic-core==2.41.5 +pydantic-core==2.46.4 # via pydantic pydantic-extra-types==2.11.1 # via @@ -992,7 +992,7 @@ regex==2026.2.28 # sacrebleu # tiktoken # transformers -requests==2.32.5 +requests==2.34.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1041,21 +1041,21 @@ rpds-py==0.30.0 # via # jsonschema # referencing -runai-model-streamer==0.15.7 +runai-model-streamer==0.16.1 # via # -c requirements/rocm.txt # -r requirements/test/rocm.in -runai-model-streamer-azure==0.15.7 +runai-model-streamer-azure==0.16.1 # via runai-model-streamer -runai-model-streamer-gcs==0.15.7 +runai-model-streamer-gcs==0.16.1 # via runai-model-streamer -runai-model-streamer-s3==0.15.7 +runai-model-streamer-s3==0.16.1 # via runai-model-streamer s3transfer==0.16.0 # via boto3 sacrebleu==2.6.0 # via lm-eval -safetensors==0.7.0 +safetensors==0.8.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1184,7 +1184,7 @@ tenacity==9.1.4 # via # gpt-oss # lm-eval -tensorizer==2.10.1 +tensorizer==2.12.1 # via # -c requirements/rocm.txt # -r requirements/test/rocm.in @@ -1194,7 +1194,7 @@ threadpoolctl==3.6.0 # via scikit-learn tifffile==2026.3.3 # via scikit-image -tiktoken==0.12.0 +tiktoken==0.14.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1242,7 +1242,7 @@ tqdm==4.67.3 # sentence-transformers # tilelang # transformers -transformers==5.5.3 +transformers==5.15.1 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1272,10 +1272,11 @@ typer==0.24.1 # huggingface-hub # perceptron # transformers -typing-extensions==4.15.0 +typing-extensions==4.16.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # aiohttp # aiosignal # albumentations # alembic @@ -1378,8 +1379,6 @@ yarl==1.23.0 # schemathesis z3-solver==4.15.4.0 # via tilelang -zipp==3.23.0 - # via importlib-metadata # The following packages were excluded from the output: # torch diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 5581d0a079..ae93c647bb 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -8,7 +8,7 @@ accelerate==1.13.0 # via -r requirements/test/xpu.in aiohappyeyeballs==2.6.1 # via aiohttp -aiohttp==3.13.4 +aiohttp==3.14.3 # via # -c requirements/common.txt # fsspec @@ -103,11 +103,11 @@ dpcpp-cpp-rt==2025.3.2 # torch evaluate==0.4.6 # via lm-eval -fastapi==0.135.2 +fastapi==0.141.1 # via # -c requirements/common.txt # gpt-oss -filelock==3.25.2 +filelock==3.32.4 # via # -c requirements/common.txt # blobfile @@ -244,7 +244,7 @@ lazy-loader==0.5 # scikit-image librosa==0.10.2.post1 # via -r requirements/test/xpu.in -llvmlite==0.47.0 +llvmlite==0.49.0 # via numba lm-eval==0.4.11 # via -r requirements/test/xpu.in @@ -304,7 +304,7 @@ nltk==3.9.4 # via rouge-score num2words==0.5.14 # via -r requirements/test/xpu.in -numba==0.65.0 +numba==0.67.0 # via # -c requirements/xpu.txt # librosa @@ -428,7 +428,7 @@ pycparser==3.0 # via cffi pycryptodomex==3.23.0 # via blobfile -pydantic==2.12.5 +pydantic==2.13.4 # via # -c requirements/common.txt # albumentations @@ -438,7 +438,7 @@ pydantic==2.12.5 # mteb # openai-harmony # pydantic-extra-types -pydantic-core==2.41.5 +pydantic-core==2.46.4 # via pydantic pydantic-extra-types==2.11.1 # via mistral-common @@ -509,7 +509,7 @@ regex==2026.3.32 # sacrebleu # tiktoken # transformers -requests==2.33.1 +requests==2.34.2 # via # -c requirements/common.txt # datasets @@ -537,7 +537,7 @@ rpds-py==0.30.0 # referencing sacrebleu==2.6.0 # via lm-eval -safetensors==0.7.0 +safetensors==0.8.0 # via # -c requirements/common.txt # accelerate @@ -634,7 +634,7 @@ threadpoolctl==3.6.0 # via scikit-learn tifffile==2026.3.3 # via scikit-image -tiktoken==0.12.0 +tiktoken==0.14.0 # via # -c requirements/common.txt # gpt-oss @@ -668,7 +668,7 @@ tqdm==4.67.3 # pqdm # sentence-transformers # transformers -transformers==5.5.3 +transformers==5.15.1 # via # -c requirements/common.txt # sentence-transformers @@ -683,9 +683,10 @@ typer==0.24.1 # via # huggingface-hub # transformers -typing-extensions==4.15.0 +typing-extensions==4.16.0 # via # -c requirements/common.txt + # aiohttp # aiosignal # albumentations # anyio diff --git a/requirements/tpu.txt b/requirements/tpu.txt index 539f2320ba..ab8fa029d8 100644 --- a/requirements/tpu.txt +++ b/requirements/tpu.txt @@ -12,4 +12,4 @@ ray[data] setuptools==78.1.0 setuptools-rust>=1.9.0 nixl==0.3.0 -tpu-inference==0.20.0 +tpu-inference==0.27.0 diff --git a/requirements/xpu.txt b/requirements/xpu.txt index 897e2080da..b7d9b5ce71 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -10,7 +10,7 @@ setuptools>=77.0.3,<81.0.0 wheel jinja2>=3.1.6 datasets # for benchmark scripts -numba == 0.65.0 # Required for N-gram speculative decoding +numba == 0.67.0 # Required for N-gram speculative decoding --extra-index-url=https://download.pytorch.org/whl/xpu torch==2.11.0+xpu torchaudio diff --git a/setup.py b/setup.py index f816f46dea..a9c6a90ab7 100644 --- a/setup.py +++ b/setup.py @@ -1392,10 +1392,10 @@ def _read_requirements(filename: str) -> list[str]: "zentorch-weekly==5.2.1.dev20260408" ], # Zentorch has weekly releases. This pulls the known-good version. "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"], - "tensorizer": ["tensorizer==2.10.1"], - "fastsafetensors": ["fastsafetensors >= 0.2.2"], + "tensorizer": ["tensorizer==2.12.1"], + "fastsafetensors": ["fastsafetensors >= 0.3.3"], "instanttensor": ["instanttensor >= 0.1.5"], - "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"], + "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.16.1"], "audio": [ "av", "scipy", @@ -1408,7 +1408,7 @@ def _read_requirements(filename: str) -> list[str]: # NOTE: When updating helion version, also update CI files: # - .buildkite/test_areas/kernels.yaml # - .buildkite/test-amd.yaml - "helion": ["helion==1.0.0"], + "helion": ["helion==1.4.0"], # Optional deps for gRPC server (vllm serve --grpc) "grpc": ["smg-grpc-servicer[vllm] >= 0.5.2"], # Optional deps for OpenTelemetry tracing diff --git a/tests/benchmarks/test_benchmark_sm70_decode.py b/tests/benchmarks/test_benchmark_sm70_decode.py index b9fe6ebb78..f5bcb3509d 100644 --- a/tests/benchmarks/test_benchmark_sm70_decode.py +++ b/tests/benchmarks/test_benchmark_sm70_decode.py @@ -3,7 +3,7 @@ import types -from benchmarks import benchmark_sm70_decode +from benchmarks import benchmark_sm70_decode, benchmark_sm70_model_tokens def test_sm70_fa2_d256_prefill_status_reports_import_error(monkeypatch): @@ -46,3 +46,61 @@ def test_sm70_fa2_d256_prefill_status_requires_dense_and_paged_ops(monkeypatch): assert status["error"] is None assert all(status["required_ops"].values()) assert all(status["optional_ops"].values()) + + +def test_sm70_fa2_d256_prefill_status_accepts_explicit_sidecar(monkeypatch): + namespace = types.SimpleNamespace() + loaded: list[str] = [] + + def load_library(path: str) -> None: + loaded.append(path) + namespace.sm70_d256_splitd_n32_dense_fwd = object() + namespace.sm70_d256_splitd_n32_paged_fwd = object() + + fake_ops = types.SimpleNamespace( + _vllm_fa2_C=namespace, + load_library=load_library, + ) + fake_torch = types.SimpleNamespace(ops=fake_ops) + monkeypatch.setattr( + benchmark_sm70_decode.importlib, + "import_module", + lambda _name: types.SimpleNamespace(), + ) + monkeypatch.setenv("VLLM_SM70_FA2_D256_LIBRARY", "/tmp/stable-fa2.so") + + status = benchmark_sm70_decode._sm70_fa2_d256_prefill_status(fake_torch) + + assert status["available"] is True + assert status["extension_file"] == "/tmp/stable-fa2.so" + assert loaded == ["/tmp/stable-fa2.so"] + + +def test_spec_decoding_delta_reports_one_sequential_prompt(): + before = [ + {"name": "vllm:spec_decode_num_drafts", "value": 10}, + {"name": "vllm:spec_decode_num_draft_tokens", "value": 30}, + {"name": "vllm:spec_decode_num_accepted_tokens", "value": 20}, + { + "name": "vllm:spec_decode_num_accepted_tokens_per_pos", + "values": [8, 7, 5], + }, + ] + after = [ + {"name": "vllm:spec_decode_num_drafts", "value": 14}, + {"name": "vllm:spec_decode_num_draft_tokens", "value": 42}, + {"name": "vllm:spec_decode_num_accepted_tokens", "value": 27}, + { + "name": "vllm:spec_decode_num_accepted_tokens_per_pos", + "values": [11, 9, 7], + }, + ] + + metrics = benchmark_sm70_model_tokens._spec_decoding_delta(before, after) + + assert metrics is not None + assert metrics["num_drafts"] == 4 + assert metrics["num_draft_tokens"] == 12 + assert metrics["num_accepted_tokens"] == 7 + assert metrics["accepted_tokens_per_pos"] == [3, 2, 2] + assert metrics["mean_acceptance_length"] == 2.75 diff --git a/tests/benchmarks/test_benchmark_sm70_dflash2_gsm8k.py b/tests/benchmarks/test_benchmark_sm70_dflash2_gsm8k.py new file mode 100644 index 0000000000..3f2b13ce93 --- /dev/null +++ b/tests/benchmarks/test_benchmark_sm70_dflash2_gsm8k.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib +from argparse import Namespace +from pathlib import Path + +import pytest + + +@pytest.fixture +def request_seeds_parser(monkeypatch): + benchmark_dir = Path(__file__).resolve().parents[2] / "benchmarks" + monkeypatch.syspath_prepend(str(benchmark_dir)) + module = importlib.import_module("benchmark_sm70_dflash2_gsm8k") + return module._request_seeds + + +@pytest.mark.parametrize( + ("request_seed", "request_seeds", "expected"), + [ + (0, None, [0]), + (-1, None, [None]), + (0, "11, 22,33", [11, 22, 33]), + ], +) +def test_request_seeds_parsing( + request_seeds_parser, request_seed, request_seeds, expected +): + args = Namespace(request_seed=request_seed, request_seeds=request_seeds) + + assert request_seeds_parser(args) == expected + + +@pytest.mark.parametrize( + ("request_seed", "request_seeds"), + [ + (-2, None), + (0, ""), + (0, "1,-2"), + (0, "7,7"), + (0, "1,invalid"), + ], +) +def test_request_seeds_rejects_ambiguous_contracts( + request_seeds_parser, request_seed, request_seeds +): + args = Namespace(request_seed=request_seed, request_seeds=request_seeds) + + with pytest.raises(ValueError): + request_seeds_parser(args) diff --git a/tests/benchmarks/test_compare_sm70_dflash2_quality_matrix.py b/tests/benchmarks/test_compare_sm70_dflash2_quality_matrix.py new file mode 100644 index 0000000000..449d7f665d --- /dev/null +++ b/tests/benchmarks/test_compare_sm70_dflash2_quality_matrix.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from benchmarks.compare_sm70_dflash2_quality_matrix import ( + _mcnemar_exact, + _ordered_arm_names, + _paired_bootstrap_delta, +) + + +def test_mcnemar_exact_counts_discordant_pairs(): + result = _mcnemar_exact( + [True, True, False, False], + [False, True, True, False], + ) + + assert result == { + "left_only": 1, + "right_only": 1, + "discordant": 2, + "p_value": 1.0, + } + + +def test_paired_bootstrap_delta_is_reproducible(): + pairs = [(1.0, 2.0), (3.0, 3.0), (4.0, 3.0)] + + left = _paired_bootstrap_delta(pairs, samples=1000, seed=7) + right = _paired_bootstrap_delta(pairs, samples=1000, seed=7) + + assert left == right + assert left["count"] == 3 + assert left["mean_delta"] == 0.0 + + +def test_matrix_arm_order_tracks_acceleration_layers(): + arms: dict[str, object] = { + "d3": {}, + "d1": {}, + "t0": {}, + "d0": {}, + "d2": {}, + } + + assert _ordered_arm_names(arms) == ["t0", "d0", "d1", "d2", "d3"] diff --git a/tests/conftest.py b/tests/conftest.py index 3eaebc38bc..882138f911 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -468,7 +468,7 @@ def _init( model_name, revision=revision, device=self.device, - automodel_args=model_kwargs, + model_kwargs=model_kwargs, trust_remote_code=trust_remote_code, ) else: diff --git a/tests/distributed/test_custom_all_reduce_dispatch.py b/tests/distributed/test_custom_all_reduce_dispatch.py new file mode 100644 index 0000000000..3703931fd3 --- /dev/null +++ b/tests/distributed/test_custom_all_reduce_dispatch.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + +def _mock_communicator() -> CustomAllreduce: + communicator = object.__new__(CustomAllreduce) + communicator.disabled = False + communicator._ptr = 0 + communicator.world_size = 4 + communicator.fully_connected = True + communicator.tp8_hierarchical = False + communicator.dispatch_max_size = 1024 * 1024 + return communicator + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_should_custom_ar_accepts_supported_dtype(dtype: torch.dtype) -> None: + communicator = _mock_communicator() + + assert communicator.should_custom_ar(torch.empty(16, dtype=dtype)) + + +@pytest.mark.parametrize( + "dtype", + [torch.float64, torch.int64, torch.int32, torch.int8, torch.uint8, torch.bool], +) +def test_should_custom_ar_rejects_unsupported_dtype(dtype: torch.dtype) -> None: + communicator = _mock_communicator() + + assert not communicator.should_custom_ar(torch.empty(16, dtype=dtype)) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 541f3d5d42..4cf07e5599 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -293,7 +293,7 @@ def test_sm70_mtp_defaults_require_env_opt_in(monkeypatch): "method": "mtp", "num_speculative_tokens": 4, "use_local_argmax_reduction": True, - "draft_sample_method": "probabilistic", + "draft_sample_method": "greedy", "attention_backend": "FLASH_ATTN_V100", } assert args.enable_prefix_caching is True @@ -379,7 +379,7 @@ def test_sm70_explicit_mtp_still_gets_safe_defaults(monkeypatch): assert args.speculative_config == { "method": "mtp", "num_speculative_tokens": 2, - "draft_sample_method": "probabilistic", + "draft_sample_method": "greedy", "use_local_argmax_reduction": True, "attention_backend": "FLASH_ATTN_V100", } diff --git a/tests/kernels/attention/test_sm70_flash_v100_dflash2_swa.py b/tests/kernels/attention/test_sm70_flash_v100_dflash2_swa.py new file mode 100644 index 0000000000..db9a2e7f82 --- /dev/null +++ b/tests/kernels/attention/test_sm70_flash_v100_dflash2_swa.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Correctness tests for the DFlash2 draft sliding-window attention shape.""" + +from __future__ import annotations + +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("seq_len", [4096, 4111]) +@torch.inference_mode() +def test_sm70_dflash2_noncausal_swa_matches_reference(seq_len: int) -> None: + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("FlashAttention-V100 regression is SM70/V100 only") + + flash_attn_v100 = pytest.importorskip("flash_attn_v100") + torch.manual_seed(0) + + q_len = 8 + num_query_heads = 8 + num_kv_heads = 2 + head_dim = 128 + block_size = 16 + window = 2048 + num_blocks = (seq_len + block_size - 1) // block_size + query = torch.randn( + 1, + q_len, + num_query_heads, + head_dim, + device="cuda", + dtype=torch.float16, + ) + key_cache = torch.randn( + num_blocks, + block_size, + num_kv_heads, + head_dim, + device="cuda", + dtype=torch.float16, + ) + value_cache = torch.randn_like(key_cache) + block_table = torch.arange(num_blocks, device="cuda", dtype=torch.int32).unsqueeze( + 0 + ) + seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32) + + actual = flash_attn_v100.flash_attn_prefill_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + causal=False, + window_size=(window, window), + ) + + key = key_cache.reshape(-1, num_kv_heads, head_dim)[:seq_len] + value = value_cache.reshape(-1, num_kv_heads, head_dim)[:seq_len] + key = key.repeat_interleave(num_query_heads // num_kv_heads, dim=1) + value = value.repeat_interleave(num_query_heads // num_kv_heads, dim=1) + query_positions = torch.arange(seq_len - q_len, seq_len, device="cuda") + key_positions = torch.arange(seq_len, device="cuda") + mask = (key_positions[None, :] >= query_positions[:, None] - window) & ( + key_positions[None, :] <= query_positions[:, None] + window + ) + scores = torch.einsum("qhd,khd->hqk", query[0].float(), key.float()) + scores *= head_dim**-0.5 + scores.masked_fill_(~mask.unsqueeze(0), float("-inf")) + probs = torch.softmax(scores, dim=-1) + reference = torch.einsum("hqk,khd->qhd", probs, value.float()) + + torch.testing.assert_close(actual[0].float(), reference, atol=2e-2, rtol=2e-2) diff --git a/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py b/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py index 84da54590b..c895d4f040 100644 --- a/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py +++ b/tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py @@ -56,6 +56,29 @@ def _make_case( return query, key_cache, value_cache, block_table, seq_lens +def _make_interleaved_case( + *, + page_size: int, + query_len: int, + prefix_len: int, +) -> tuple[torch.Tensor, ...]: + total_len = prefix_len + query_len + logical_pages = math.ceil(total_len / page_size) + physical_pages = logical_pages + 3 + source_shape = (physical_pages, 2, page_size, 1, 256) + source = torch.randn(source_shape, dtype=torch.float16, device="cuda").mul_(0.25) + cache = source.to(torch.float8_e5m2).view(torch.uint8) + key_cache, value_cache = cache.unbind(1) + block_table = torch.randperm(physical_pages, dtype=torch.int32, device="cuda")[ + :logical_pages + ].view(1, -1) + query = torch.randn((query_len, 6, 256), dtype=torch.float16, device="cuda").mul_( + 0.25 + ) + seq_lens = torch.tensor([total_len], dtype=torch.int32, device="cuda") + return query, key_cache, value_cache, block_table, seq_lens + + def _reference( query: torch.Tensor, key_cache: torch.Tensor, @@ -104,10 +127,17 @@ def _reference( (16, 5, 507, False), (16, 6, 506, False), (16, 8, 1016, False), + (16, 16, 1008, True), (16, 3, 1025, True), (784, 5, 2049, True), (1648, 8, 4097, True), + (1648, 16, 4089, True), + (1728, 8, 4097, True), + (1728, 16, 4089, True), (3296, 8, 8193, True), + (3296, 16, 8185, True), + (3456, 8, 8193, True), + (3456, 16, 8185, True), ], ) @torch.inference_mode() @@ -194,13 +224,16 @@ def test_grouped_verify_matches_fp32_reference_with_random_pages( torch.testing.assert_close(one_pass, two_pass, atol=6.2e-5, rtol=2.0e-3) +@pytest.mark.parametrize("query_len", [8, 16]) @torch.inference_mode() -def test_grouped_verify_cuda_graph_replay_tracks_runtime_seq_len() -> None: +def test_grouped_verify_cuda_graph_replay_tracks_runtime_seq_len( + query_len: int, +) -> None: flash_attn_v100 = _require_grouped_verify() torch.manual_seed(20260825) query, key_cache, value_cache, block_table, seq_lens = _make_case( page_size=1648, - query_len=8, + query_len=query_len, prefix_len=4097, ) output = torch.empty_like(query) @@ -243,3 +276,53 @@ def test_grouped_verify_cuda_graph_replay_tracks_runtime_seq_len() -> None: difference = output.float().sub(expected.float()).abs() assert difference.max().item() <= 6.2e-5 assert difference.mean().item() <= 6.0e-6 + + +@pytest.mark.parametrize("page_size", [1648, 3296]) +@pytest.mark.parametrize("query_len", [8, 16]) +@torch.inference_mode() +def test_grouped_verify_fixed_interleaved_is_bitwise( + page_size: int, query_len: int, monkeypatch: pytest.MonkeyPatch +) -> None: + flash_attn_v100 = _require_grouped_verify() + torch.manual_seed(20260826 + page_size) + query, key_cache, value_cache, block_table, seq_lens = _make_interleaved_case( + page_size=page_size, + query_len=query_len, + prefix_len=8193, + ) + + monkeypatch.setenv("VLLM_FLASH_V100_DFLASH2_FIXED_INTERLEAVED", "0") + control = flash_attn_v100.flash_attn_grouped_verify_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + one_pass=True, + ).clone() + + monkeypatch.setenv("VLLM_FLASH_V100_DFLASH2_FIXED_INTERLEAVED", "1") + monkeypatch.setenv("VLLM_FLASH_V100_DFLASH2_STAGE_PAGE_IDS", "0") + fixed = flash_attn_v100.flash_attn_grouped_verify_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + one_pass=True, + ).clone() + + monkeypatch.setenv("VLLM_FLASH_V100_DFLASH2_STAGE_PAGE_IDS", "1") + staged = flash_attn_v100.flash_attn_grouped_verify_paged( + query, + key_cache, + value_cache, + block_table, + seq_lens, + one_pass=True, + ).clone() + torch.accelerator.synchronize() + + assert torch.equal(fixed, control) + assert torch.equal(staged, control) diff --git a/tests/kernels/attention/test_sm70_flash_v100_splitkv3.py b/tests/kernels/attention/test_sm70_flash_v100_splitkv3.py index 4d829aaeb9..6e63720ed1 100644 --- a/tests/kernels/attention/test_sm70_flash_v100_splitkv3.py +++ b/tests/kernels/attention/test_sm70_flash_v100_splitkv3.py @@ -37,18 +37,14 @@ def _run_case( dtype=torch.float16, ) value_cache = torch.randn_like(key_cache) - block_table = torch.tensor( - [page_ids], device="cuda", dtype=torch.int32 - ) + block_table = torch.tensor([page_ids], device="cuda", dtype=torch.int32) sequence_lengths = torch.full( (batch_size,), sequence_len, device="cuda", dtype=torch.int32 ) softmax_scale = head_dim**-0.5 unsplit_out = torch.empty_like(query) - unsplit_lse = torch.empty( - query.shape[:-1], device="cuda", dtype=torch.float32 - ) + unsplit_lse = torch.empty(query.shape[:-1], device="cuda", dtype=torch.float32) split_out = torch.empty_like(query) split_lse = torch.empty_like(unsplit_lse) split_tmp_out = torch.empty( @@ -202,6 +198,4 @@ def test_splitkv3_python_wrapper_reuses_stream_workspace() -> None: rtol=0.0, atol=0.001953125, ) - torch.testing.assert_close( - second_lse, result["unsplit_lse"], rtol=0.0, atol=0.001 - ) + torch.testing.assert_close(second_lse, result["unsplit_lse"], rtol=0.0, atol=0.001) diff --git a/tests/kernels/moe/test_fused_topk.py b/tests/kernels/moe/test_fused_topk.py index 825cd20263..623bd3d872 100644 --- a/tests/kernels/moe/test_fused_topk.py +++ b/tests/kernels/moe/test_fused_topk.py @@ -8,10 +8,14 @@ import pytest import torch +import vllm._custom_ops as ops from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( fused_topk_bias, ) -from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk +from vllm.model_executor.layers.fused_moe.router.fused_topk_router import ( + _sm70_qwen38_router_topk, + fused_topk, +) from vllm.platforms import current_platform @@ -44,6 +48,54 @@ def torch_topk( return topk_weights, topk_ids +@pytest.mark.skipif( + not current_platform.is_device_capability(70), + reason="Qwen3.8 router specialization requires SM70", +) +@pytest.mark.parametrize( + "case", ["random", "ties", "signed_zero", "nan", "inf", "negative_inf"] +) +def test_sm70_qwen38_router_topk(case: str): + torch.manual_seed(0) + gating_output = torch.randn(1, 512, dtype=torch.float16, device="cuda") + if case == "ties": + gating_output[:, :16] = 1.0 + elif case == "signed_zero": + gating_output.zero_() + gating_output[:, 1::2] = -0.0 + elif case == "nan": + gating_output[:, 0] = float("nan") + elif case == "inf": + gating_output[:, 0] = float("inf") + elif case == "negative_inf": + gating_output.fill_(-float("inf")) + + expected_weights = torch.empty(1, 10, dtype=torch.float32, device="cuda") + expected_ids = torch.empty(1, 10, dtype=torch.int32, device="cuda") + expected_rows = torch.empty_like(expected_ids) + ops.topk_softmax( + expected_weights, + expected_ids, + expected_rows, + gating_output, + True, + ) + + actual_weights = torch.empty_like(expected_weights) + actual_ids = torch.empty_like(expected_ids) + actual_rows = torch.empty_like(expected_rows) + _sm70_qwen38_router_topk( + actual_weights, + actual_ids, + actual_rows, + gating_output, + ) + + torch.testing.assert_close(actual_ids, expected_ids, atol=0, rtol=0) + torch.testing.assert_close(actual_rows, expected_rows, atol=0, rtol=0) + torch.testing.assert_close(actual_weights, expected_weights, atol=1e-7, rtol=1e-7) + + @pytest.mark.skipif( not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." ) diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 7e8ab171bf..dc596f8795 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -474,9 +474,7 @@ def test_sm70_single_token_weighted_reduce_matches_moe_unpermute(): ) topk_weights = torch.randn((1, topk), device=device, dtype=torch.float32) dummy_hidden = torch.randn((1, hidden), device=device, dtype=torch.float16) - _, _, offsets, inv, _ = moe_permute( - dummy_hidden, None, topk_ids, n_expert - ) + _, _, offsets, inv, _ = moe_permute(dummy_hidden, None, topk_ids, n_expert) sorted_output_storage = torch.randn( (topk, hidden + pad), device=device, dtype=torch.float16 ) @@ -484,12 +482,8 @@ def test_sm70_single_token_weighted_reduce_matches_moe_unpermute(): expected = torch.empty((1, hidden), device=device, dtype=torch.float16) actual = torch.empty_like(expected) - reference_sorted_output = ( - sorted_output.contiguous() if pad else sorted_output - ) - moe_unpermute( - expected, reference_sorted_output, topk_weights, inv, offsets - ) + reference_sorted_output = sorted_output.contiguous() if pad else sorted_output + moe_unpermute(expected, reference_sorted_output, topk_weights, inv, offsets) torch.ops._C.awq_moe_single_token_weighted_reduce_out( sorted_output, topk_weights, diff --git a/tests/kernels/quantization/test_nvfp4_emulation.py b/tests/kernels/quantization/test_nvfp4_emulation.py index b84bea7214..b5099f6dcd 100644 --- a/tests/kernels/quantization/test_nvfp4_emulation.py +++ b/tests/kernels/quantization/test_nvfp4_emulation.py @@ -63,6 +63,74 @@ def test_nvfp4_sm70_e4m3_scale_fallback_matches_torch_float8(monkeypatch) -> Non torch.testing.assert_close(actual, expected, atol=0, rtol=0) +@pytest.mark.skipif( + not (current_platform.is_cuda() and current_platform.has_device_capability(70)), + reason="Integer-decoded NVFP4 weight dequantization requires CUDA SM70+.", +) +@pytest.mark.parametrize("output_dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("use_3d", [False, True]) +def test_integer_e4m3_weight_dequantization_matches_reference( + monkeypatch, output_dtype: torch.dtype, use_3d: bool +) -> None: + """The integer E4M3 decoder must bit-match the PyTorch fallback.""" + block_size = 16 + raw_scales = torch.arange(256, dtype=torch.uint8) + rows = raw_scales.numel() + tensor_sf = raw_scales.reshape(rows, 1).view(torch.float8_e4m3fn).cuda() + tensor_fp4 = ( + torch.arange(rows * (block_size // 2), dtype=torch.int32) + .remainder(256) + .to(torch.uint8) + .reshape(rows, block_size // 2) + .cuda() + ) + global_scale = torch.tensor(0.75, dtype=torch.float32, device="cuda") + finite_rows = (raw_scales & 0x7F) != 0x7F + if use_3d: + tensor_sf = tensor_sf.reshape(2, rows // 2, 1) + tensor_fp4 = tensor_fp4.reshape(2, rows // 2, block_size // 2) + global_scale = torch.tensor([0.5, 1.25], device="cuda") + finite_rows = finite_rows.reshape(2, rows // 2) + + monkeypatch.setattr( + nvfp4_emulation_utils, "_TRITON_NVFP4_EMULATION_SUPPORTED", False + ) + monkeypatch.setattr( + nvfp4_emulation_utils, + "_TRITON_NVFP4_DEQUANTIZATION_SUPPORTED", + True, + ) + actual = dequantize_to_dtype( + tensor_fp4, + tensor_sf, + global_scale, + output_dtype, + block_size, + swizzle=False, + ) + + with monkeypatch.context() as reference_patch: + reference_patch.setattr( + nvfp4_emulation_utils, + "_TRITON_NVFP4_DEQUANTIZATION_SUPPORTED", + False, + ) + expected = dequantize_to_dtype( + tensor_fp4, + tensor_sf, + global_scale, + output_dtype, + block_size, + swizzle=False, + ) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0, equal_nan=True) + assert torch.equal( + actual[finite_rows].view(torch.int16), + expected[finite_rows].view(torch.int16), + ) + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="Triton NVFP4 kernel requires CUDA.", diff --git a/tests/kernels/test_deepseek_v4_sm70_qnorm_rope_kv_insert.py b/tests/kernels/test_deepseek_v4_sm70_qnorm_rope_kv_insert.py index ef52036b1c..3cb8dc13ba 100644 --- a/tests/kernels/test_deepseek_v4_sm70_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_deepseek_v4_sm70_qnorm_rope_kv_insert.py @@ -152,7 +152,7 @@ def run(enabled: bool) -> tuple[torch.Tensor, torch.Tensor]: eps=1e-6, block_size=block_size, ) - torch.cuda.synchronize() + torch.accelerator.synchronize() return q, storage legacy_q, legacy_storage = run(False) diff --git a/tests/kernels/test_fused_recurrent_packed_decode.py b/tests/kernels/test_fused_recurrent_packed_decode.py index ce9a687606..10e6e6af42 100644 --- a/tests/kernels/test_fused_recurrent_packed_decode.py +++ b/tests/kernels/test_fused_recurrent_packed_decode.py @@ -151,9 +151,7 @@ def test_fused_recurrent_packed_decode_cuda_graph_replay_uses_runtime_indices(): mixed_qkv_static = torch.empty((B, qkv_dim), device=device, dtype=dtype) a_static = torch.empty((B, HV), device=device, dtype=dtype) b_static = torch.empty((B, HV), device=device, dtype=dtype) - state_static = torch.empty( - (num_state_slots, HV, V, K), device=device, dtype=dtype - ) + state_static = torch.empty((num_state_slots, HV, V, K), device=device, dtype=dtype) indices_static = torch.empty((B,), device=device, dtype=torch.int32) out_static = torch.empty((B, 1, HV, V), device=device, dtype=dtype) @@ -161,9 +159,7 @@ def test_fused_recurrent_packed_decode_cuda_graph_replay_uses_runtime_indices(): warm_mixed_qkv = torch.randn((B, qkv_dim), device=device, dtype=dtype) warm_a = torch.randn((B, HV), device=device, dtype=dtype) warm_b = torch.randn((B, HV), device=device, dtype=dtype) - warm_state = torch.randn( - (num_state_slots, HV, V, K), device=device, dtype=dtype - ) + warm_state = torch.randn((num_state_slots, HV, V, K), device=device, dtype=dtype) warm_indices = torch.tensor([0, 1, -1, -1], device=device, dtype=torch.int32) warm_out = torch.empty_like(out_static) fused_recurrent_gated_delta_rule_packed_decode( @@ -183,9 +179,7 @@ def test_fused_recurrent_packed_decode_cuda_graph_replay_uses_runtime_indices(): capture_mixed_qkv = torch.randn((B, qkv_dim), device=device, dtype=dtype) capture_a = torch.randn((B, HV), device=device, dtype=dtype) capture_b = torch.randn((B, HV), device=device, dtype=dtype) - capture_state = torch.randn( - (num_state_slots, HV, V, K), device=device, dtype=dtype - ) + capture_state = torch.randn((num_state_slots, HV, V, K), device=device, dtype=dtype) capture_indices = torch.tensor([6, -1, 7, -1], device=device, dtype=torch.int32) mixed_qkv_static.copy_(capture_mixed_qkv) a_static.copy_(capture_a) diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py new file mode 100644 index 0000000000..4835c79347 --- /dev/null +++ b/tests/kernels/test_kpool_decode_update_batched.py @@ -0,0 +1,452 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit test for the batched kpool decode-update kernel. + +Validates ``kpool_decode_update_and_maybe_write_cache_batched`` against an +independent pure-torch reference that replicates the per-request, in-position +order semantics: stash each token into a paged tail ring; on pool completion +(``pos % pool_size == pool_size-1``) softmax(gate+ape)-weighted sum + Hadamard-128 ++ per-vector fp8 absmax quant + write to the indexer K cache. Covers +no-completion, completion-at-end, completion-mid-batch, non-uniform padding, +plain decode, plus a randomized fuzz pass. + +The kernel iterates each request's ``next_n`` tokens in position order inside +one program (grid = num_requests) to preserve the pool-completion +read-after-stash dependency; the reference mirrors that ordering. + +``test_decode_writer_matches_prefill_writer`` is deliberately NOT +reference-based: it checks the decode writer against the *prefill* writer +(``kpool_compress_and_write_cache``), the invariant that actually matters in +production. A hand-written reference can drift to match a buggy kernel -- that +is exactly how the stash-gating bug (intra-pool tokens never entering the tail +ring, because the stash was gated on the pool-granular ``slot_mapping``) stayed +green here. +""" + +import math + +import pytest +import torch + +from vllm.models.glm5next.nvidia.ops.kpool_compress import ( + kpool_compress_and_write_cache, + kpool_decode_update_and_maybe_write_cache_batched, +) + +HEAD_DIM = 128 +POOL_SIZE = 16 +PAGE_SIZE = 64 +NUM_BLOCKS = 32 +ROUND_SCALE = True +FP8_MAX = 448.0 + + +def _make_caches(): + kv = torch.zeros( + NUM_BLOCKS, PAGE_SIZE, HEAD_DIM + 4, dtype=torch.uint8, device="cuda" + ) + tail = torch.zeros( + NUM_BLOCKS, 2, POOL_SIZE, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + return kv, tail + + +def _tail_slot_for(blocks, pos): + """tail_slot = block*POOL + pos%POOL; each request owns a distinct tail block.""" + blk = torch.tensor(blocks, device=pos.device, dtype=torch.int32).unsqueeze(1) + return (blk * POOL_SIZE + pos % POOL_SIZE).to(torch.int32) + + +def _seed_prior(tail, blocks, n_prior, seed=42): + if n_prior <= 0: + return + g = torch.Generator(device=tail.device).manual_seed(seed) + prior_k = torch.randn( + len(blocks), + n_prior, + HEAD_DIM, + dtype=torch.bfloat16, + device=tail.device, + generator=g, + ) + prior_s = torch.randn( + len(blocks), + n_prior, + HEAD_DIM, + dtype=torch.bfloat16, + device=tail.device, + generator=g, + ) + for i, blk in enumerate(blocks): + tail[blk, 0, :n_prior, :] = prior_k[i] + tail[blk, 1, :n_prior, :] = prior_s[i] + + +def _hadamard128_torch(x: torch.Tensor) -> torch.Tensor: + """Reference Hadamard-128 on the last dim (must be 128).""" + n = x.shape[-1] + assert n == 128 + h = torch.tensor([[1.0, 1.0], [1.0, -1.0]], dtype=torch.float32, device=x.device) + while h.shape[0] < n: + h = torch.cat([torch.cat([h, h], dim=1), torch.cat([h, -h], dim=1)], dim=0) + h = h / math.sqrt(n) + return x @ h + + +def _torch_reference( + kv: torch.Tensor, + tail: torch.Tensor, + tail_slot: torch.Tensor, + key: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + slot_map: torch.Tensor, + pos: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Independent reference for the batched decode-update kernel. + + For each request, iterate its next_n tokens in order. On each token: + - if pos%POOL == POOL-1 and pos_valid: compress the pool (slots + [pool_start..pool_start+POOL-1], current token via is_current) and write + fp8 K + fp32 scale to kv_cache at cache_loc. + - always stash the current token's K/score into tail[block, pos%POOL]. + """ + kv = kv.clone() + tail = tail.clone() + B, next_n = pos.shape + # The indexer K cache is [num_blocks, PAGE_SIZE, HEAD_DIM+4] uint8 but the + # kernels interpret each page as [HEAD_DIM*PAGE_SIZE bytes of K (token-major) + # | 4*PAGE_SIZE bytes of fp32 scale (token-major)]. Operate on a flat byte + # view so the reference writes K and scale at the exact offsets the kernel + # uses (page_base + tok*HEAD_DIM for K; page_base + HEAD_DIM*PAGE_SIZE + + # tok*4 for scale). + page_bytes = PAGE_SIZE * (HEAD_DIM + 4) + k_region = HEAD_DIM * PAGE_SIZE + tail_slot_cpu = tail_slot.cpu().tolist() + slot_map_cpu = slot_map.cpu().tolist() + pos_cpu = pos.cpu().tolist() + key_cpu = key.float().cpu() + score_cpu = score.float().cpu() + ape_cpu = ape.cpu() + tail_cpu = tail.float().cpu() + kv_flat = kv.view(torch.uint8).reshape(-1).cpu() + + for b in range(B): + for t in range(next_n): + cache_loc = slot_map_cpu[b][t] + p = pos_cpu[b][t] + pos_valid = cache_loc >= 0 and p >= 0 + safe_pos = max(p, 0) + slot = safe_pos % POOL_SIZE + phys_slot = safe_pos % POOL_SIZE + # Per-token block derivation (a leading invalid sentinel must not + # poison the base for the rest of the request); clamped like the + # kernel so an invalid entry can't form a negative base. + block = max(tail_slot_cpu[b][t], 0) // POOL_SIZE + + cur_key = key_cpu[b, t] + cur_score = score_cpu[b, t] + + if pos_valid and slot == POOL_SIZE - 1: + pool_logical_start = safe_pos - slot + pool_scores = [] + pool_ks = [] + for ps in range(POOL_SIZE): + is_current = ps == slot + phys = (pool_logical_start + ps) % POOL_SIZE + if is_current: + s = cur_score + k = cur_key + else: + s = tail_cpu[block, 1, phys] + k = tail_cpu[block, 0, phys] + s = s + ape_cpu[ps] + pool_scores.append(s) + pool_ks.append(k) + pool_scores = torch.stack(pool_scores) # [POOL, D] + pool_ks = torch.stack(pool_ks) # [POOL, D] + max_score = pool_scores.max(dim=0).values + prob = torch.exp(pool_scores - max_score) + denom = prob.sum(dim=0) + acc = (pool_ks * prob).sum(dim=0) + x = (acc / denom).to(torch.bfloat16).to(torch.float32) + x = _hadamard128_torch(x).to(torch.bfloat16).to(torch.float32) + absmax = torch.clamp(x.abs().max(), min=1e-4) + if ROUND_SCALE: + scale = torch.exp2(torch.ceil(torch.log2(absmax / FP8_MAX))) + else: + scale = absmax / FP8_MAX + quantized = torch.clamp(x / scale, -FP8_MAX, FP8_MAX).to( + torch.float8_e4m3fn + ) + # write K and scale at the separated-layout offsets + loc = cache_loc + loc_page_index = loc // PAGE_SIZE + loc_tok = loc % PAGE_SIZE + page_base = loc_page_index * page_bytes + k_off = page_base + loc_tok * HEAD_DIM + s_off = page_base + k_region + loc_tok * 4 + kv_flat[k_off : k_off + HEAD_DIM] = quantized.view(torch.uint8) + kv_flat[s_off : s_off + 4] = scale.detach().reshape(1).view(torch.uint8) + + # stash -- gated on the TOKEN-granular tail slot, not on pos_valid. + # pos_valid keys off the pool-granular cache_loc, which is -1 for + # every token that is not the pool's last, so gating the stash on it + # would drop all intra-pool tokens. + if p >= 0 and tail_slot_cpu[b][t] >= 0: + tail_cpu[block, 0, phys_slot] = cur_key + tail_cpu[block, 1, phys_slot] = cur_score + + kv_out = kv_flat.view(NUM_BLOCKS, PAGE_SIZE, HEAD_DIM + 4).to(device="cuda") + return kv_out, tail_cpu.to(torch.bfloat16).to(device="cuda") + + +def _assert_eq(r_ref, r_kern): + kv_ref, tail_ref = r_ref + kv_kern, tail_kern = r_kern + assert torch.equal(kv_ref, kv_kern), ( + "kv_cache differs: max diff " + f"{(kv_ref.int() - kv_kern.int()).abs().max().item()}" + ) + assert torch.equal(tail_ref, tail_kern), ( + "tail_kv_cache differs: max diff " + f"{(tail_ref.float() - tail_kern.float()).abs().max().item()}" + ) + + +def _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos): + kv = kv.clone() + tail = tail.clone() + kpool_decode_update_and_maybe_write_cache_batched( + kv, + tail, + tail_slot, + key, + score, + ape, + slot_map, + pos, + POOL_SIZE, + HEAD_DIM, + round_scale=ROUND_SCALE, + ) + return kv, tail + + +@pytest.mark.parametrize("pool_size", [4, 16]) +def test_decode_writer_matches_prefill_writer(pool_size): + """Compare production decode and prefill writers for pool sizes 4 and 16.""" + n_pools, page, nblk = 8, 64, 4 + n_tok = n_pools * pool_size + dev = "cuda" + torch.manual_seed(0) + k = torch.randn(n_tok, HEAD_DIM, dtype=torch.bfloat16, device=dev) + score = torch.randn(n_tok, HEAD_DIM, dtype=torch.bfloat16, device=dev) + ape = torch.randn(pool_size, HEAD_DIM, dtype=torch.float32, device=dev) + + kv_prefill = torch.zeros(nblk, page, HEAD_DIM + 4, dtype=torch.uint8, device=dev) + kpool_compress_and_write_cache( + kv_prefill, + k.view(n_pools, pool_size, HEAD_DIM), + score.view(n_pools, pool_size, HEAD_DIM), + ape, + torch.arange(n_pools, dtype=torch.int64, device=dev), + pool_size=pool_size, + head_dim=HEAD_DIM, + round_scale=ROUND_SCALE, + ) + + # One request owning tail block 0, fed one token per decode step. + kv_decode = torch.zeros_like(kv_prefill) + tail = torch.zeros(nblk, 2, pool_size, HEAD_DIM, dtype=torch.bfloat16, device=dev) + for t in range(n_tok): + completes = t % pool_size == pool_size - 1 + kpool_decode_update_and_maybe_write_cache_batched( + kv_decode, + tail, + # token-granular: every token has a valid tail slot + torch.tensor([[t % pool_size]], dtype=torch.int32, device=dev), + k[t].view(1, 1, HEAD_DIM), + score[t].view(1, 1, HEAD_DIM), + ape, + # pool-granular: only the pool's last token carries a cache slot + torch.tensor( + [[t // pool_size if completes else -1]], dtype=torch.int32, device=dev + ), + torch.tensor([[t]], dtype=torch.int32, device=dev), + pool_size, + HEAD_DIM, + round_scale=ROUND_SCALE, + ) + + differing = [ + p + for p in range(n_pools) + if not torch.equal( + kv_prefill[p // page, p % page], kv_decode[p // page, p % page] + ) + ] + assert not differing, ( + f"decode-written pools differ from prefill-written pools: " + f"{len(differing)}/{n_pools} (pool_size={pool_size}, first={differing[:5]})" + ) + + +def test_leading_invalid_tail_slot(): + """A request whose FIRST token carries an invalid (-1) tail slot while a + later token is a real pool completion. + + The tail block must be derived per token, not from token 0: a leading + invalid sentinel would otherwise poison the base address for the whole + request (out-of-bounds tail reads on the completion). + """ + torch.manual_seed(0) + B, next_n, blocks = 2, 4, [3, 5] + # req 0: token 0 invalid (pos -1), tokens 1..3 valid, completion at pos 15 + # req 1: all valid, no completion + pos = torch.tensor( + [[-1, 13, 14, 15], [4, 5, 6, 7]], dtype=torch.int32, device="cuda" + ) + safe_pos = torch.where(pos >= 0, pos, 0) + tail_slot = _tail_slot_for(blocks, safe_pos) + # leading invalid entry carries the -1 sentinel, as the scatter path emits + tail_slot[0, 0] = -1 + slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") + slot_map[0, 3] = 15 # req 0 completes its pool on the last verify token + + key = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + score = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + ape = torch.randn(POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda") + + kv, tail = _make_caches() + _seed_prior(tail, blocks, 13) + r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos) + r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos) + _assert_eq(r_ref, r_kern) + + +@pytest.mark.parametrize( + "case_id", + [ + "no_completion", + "completion_at_end", + "completion_mid_batch", + "non_uniform_padding", + "plain_decode", + ], +) +def test_batched_matches_reference(case_id): + torch.manual_seed(0) + if case_id == "no_completion": + B, next_n, blocks = 3, 4, [0, 1, 2] + pos = ( + torch.arange(next_n, device="cuda", dtype=torch.int32) + .unsqueeze(0) + .expand(B, -1) + .contiguous() + ) + slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") + n_prior = 0 + elif case_id == "completion_at_end": + B, next_n, blocks = 2, 4, [0, 1] + pos = torch.tensor( + [[12, 13, 14, 15], [12, 13, 14, 15]], dtype=torch.int32, device="cuda" + ) + slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") + slot_map[:, 3] = torch.tensor( + [15, PAGE_SIZE + 15], dtype=torch.int32, device="cuda" + ) + n_prior = POOL_SIZE - next_n + elif case_id == "completion_mid_batch": + B, next_n, blocks = 3, 4, [0, 1, 2] + pos = torch.tensor([[13, 14, 15, 16]] * B, dtype=torch.int32, device="cuda") + slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") + slot_map[:, 2] = torch.tensor( + [15, PAGE_SIZE + 15, 2 * PAGE_SIZE + 15], dtype=torch.int32, device="cuda" + ) + n_prior = 13 + elif case_id == "non_uniform_padding": + B, next_n, blocks = 2, 4, [0, 1] + pos = torch.tensor( + [[12, 13, 14, 15], [12, 13, -1, -1]], dtype=torch.int32, device="cuda" + ) + slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") + slot_map[0, 3] = 15 + n_prior = POOL_SIZE - 4 + else: # plain_decode + B, next_n, blocks = 4, 1, [0, 1, 2, 3] + pos = torch.tensor([[5], [6], [7], [8]], dtype=torch.int32, device="cuda") + slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") + n_prior = 0 + + if case_id == "non_uniform_padding": + safe_pos = torch.where(pos >= 0, pos, 0) + tail_slot = torch.where(pos >= 0, _tail_slot_for(blocks, safe_pos), 0) + else: + tail_slot = _tail_slot_for(blocks, pos) + + key = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + score = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + ape = torch.randn(POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda") + + kv, tail = _make_caches() + _seed_prior(tail, blocks, n_prior) + r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos) + r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos) + _assert_eq(r_ref, r_kern) + + +@pytest.mark.parametrize("seed", list(range(20))) +def test_batched_matches_reference_fuzz(seed): + """Random B / next_n / start positions; covers 0, 1, and multi completion.""" + g = torch.Generator(device="cuda").manual_seed(seed) + B = int(torch.randint(1, 6, (1,), generator=g, device="cuda").item()) + next_n = int(torch.randint(1, 8, (1,), generator=g, device="cuda").item()) + blocks = list(range(B)) + + starts = torch.randint(0, 33, (B,), generator=g, device="cuda", dtype=torch.int32) + pos = starts.unsqueeze(1) + torch.arange( + next_n, device="cuda", dtype=torch.int32 + ).unsqueeze(0) + tail_slot = _tail_slot_for(blocks, pos) + + is_completion = pos % POOL_SIZE == POOL_SIZE - 1 + blk = torch.tensor(blocks, device="cuda", dtype=torch.int32).unsqueeze(1) + pool_slot = blk * PAGE_SIZE + (POOL_SIZE - 1) + slot_map = torch.where(is_completion, pool_slot, torch.full_like(pos, -1)) + + key = torch.randn( + B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda", generator=g + ) + score = torch.randn( + B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda", generator=g + ) + ape = torch.randn( + POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda", generator=g + ) + + kv, tail = _make_caches() + prior_g = torch.Generator(device="cuda").manual_seed(seed + 1000) + for b in range(B): + n_prior = int(starts[b].item()) % POOL_SIZE + if n_prior > 0: + pk = torch.randn( + n_prior, + HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + generator=prior_g, + ) + ps = torch.randn( + n_prior, + HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + generator=prior_g, + ) + tail[blocks[b], 0, :n_prior, :] = pk + tail[blocks[b], 1, :n_prior, :] = ps + + r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos) + r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos) + _assert_eq(r_ref, r_kern) diff --git a/tests/kernels/test_mhc_sm70_fp16.py b/tests/kernels/test_mhc_sm70_fp16.py index a0086f77e7..0594841fda 100644 --- a/tests/kernels/test_mhc_sm70_fp16.py +++ b/tests/kernels/test_mhc_sm70_fp16.py @@ -4,6 +4,13 @@ import torch from vllm.model_executor.kernels.mhc import tilelang as mhc_tilelang +from vllm.model_executor.kernels.mhc.torch import ( + mhc_fused_post_pre_torch, + mhc_post_torch, + mhc_pre_torch, +) +from vllm.model_executor.kernels.mhc.triton import sm70_mhc_pre_norm_from_staging +from vllm.model_executor.layers import mhc as mhc_layer from vllm.platforms.interface import DeviceCapability from vllm.utils.import_utils import has_tilelang @@ -71,6 +78,21 @@ def test_mhc_fake_paths_preserve_fp16_graph_metadata() -> None: pre = torch.ops.vllm.mhc_pre_tilelang( residual, fn, scale, base, 1e-6, 1e-6, 1e-6, 1.0, 2 ) + broadcast = torch.ops.vllm.mhc_pre_broadcast_tilelang( + x, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 2, + 1, + torch.empty((128,), dtype=torch.float16, device="meta"), + 1e-6, + torch.empty((24, 128), dtype=torch.float32, device="meta"), + ) fused = torch.ops.vllm.mhc_fused_post_pre_tilelang( x, residual, post, comb, fn, scale, base, 1e-6, 1e-6, 1e-6, 1.0, 2 ) @@ -85,6 +107,8 @@ def test_mhc_fake_paths_preserve_fp16_graph_metadata() -> None: assert pre[2].dtype == torch.float16 assert pre[2].shape == (2, 128) + assert broadcast[0].shape == (2, 4, 128) + assert broadcast[3].dtype == torch.float16 assert fused[0].dtype == torch.float16 assert fused[3].dtype == torch.float16 assert fused[3].shape == (2, 128) @@ -115,11 +139,401 @@ def run( assert output[2].shape == (2, 128) +def test_mhc_broadcast_fake_path_compiles_fullgraph() -> None: + def run( + residual: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: torch.Tensor, + fn_broadcast: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.ops.vllm.mhc_pre_broadcast_tilelang( + residual, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 2, + 1, + norm_weight, + 1e-6, + fn_broadcast, + ) + + compiled = torch.compile(run, backend="eager", fullgraph=True) + output = compiled( + torch.empty((2, 128), dtype=torch.float16, device="meta"), + torch.empty((24, 512), dtype=torch.float32, device="meta"), + torch.empty((3,), dtype=torch.float32, device="meta"), + torch.empty((24,), dtype=torch.float32, device="meta"), + torch.empty((128,), dtype=torch.float16, device="meta"), + torch.empty((24, 128), dtype=torch.float32, device="meta"), + ) + + assert output[0].shape == (2, 4, 128) + assert output[3].dtype == torch.float16 + + +def test_mhc_torch_fp16_fused_prefill_matches_fp32_reference() -> None: + torch.manual_seed(17) + num_tokens = 17 + hidden_size = 16 + hc_mult = 4 + hc_mult3 = 24 + residual = torch.randn(num_tokens, hc_mult, hidden_size, dtype=torch.float16) + x = torch.randn(num_tokens, hidden_size, dtype=torch.float16) + post = torch.randn(num_tokens, hc_mult, 1, dtype=torch.float32) + comb = torch.randn(num_tokens, hc_mult, hc_mult, dtype=torch.float32) + fn = torch.randn(hc_mult3, hc_mult * hidden_size, dtype=torch.float32) + scale = torch.tensor([0.5, 0.75, 0.25], dtype=torch.float32) + base = torch.randn(hc_mult3, dtype=torch.float32) + norm_weight = torch.randn(hidden_size, dtype=torch.float16) + + output = mhc_fused_post_pre_torch( + x, + residual, + post, + comb, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 3, + norm_weight=norm_weight, + norm_eps=1e-6, + ) + residual_ref = ( + torch.einsum("bij,bih->bjh", comb, residual.float()) + + post * x.float().unsqueeze(1) + ).half() + pre_ref = mhc_pre_torch( + residual_ref, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 3, + norm_weight=norm_weight, + norm_eps=1e-6, + ) + + torch.testing.assert_close(output[0], residual_ref, rtol=0, atol=0) + for actual, expected in zip(output[1:], pre_ref): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert output[-1].dtype == torch.float16 + + +def test_mhc_sm70_fallback_policy(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(mhc_layer, "current_platform", _FakeCudaPlatform(False)) + decode = torch.empty((1, 4, 4096), dtype=torch.float16) + prefill = torch.empty((17, 4, 4096), dtype=torch.float16) + unsupported = torch.empty((1, 4, 8), dtype=torch.float16) + + assert not mhc_layer._use_sm70_fp16_mhc_fallback(decode, prefill_only=True) + assert not mhc_layer._use_sm70_fp16_mhc_fallback(prefill, prefill_only=True) + assert not mhc_layer._use_sm70_fp16_mhc_fallback(decode, prefill_only=False) + assert mhc_layer._use_sm70_fp16_mhc_fallback(unsupported, prefill_only=False) + + +@pytest.mark.skipif( + not ( + mhc_tilelang.current_platform.is_cuda() + and mhc_tilelang.current_platform.get_device_capability() + == DeviceCapability(7, 0) + and hasattr(torch.ops._C, "sm70_glm_mhc_pre_norm_out") + ), + reason="native NVIDIA V100/SM70 mHC CUDA op required", +) +def test_mhc_sm70_native_final_stage_graph_matches_eager_bitwise() -> None: + torch.manual_seed(20260831) + device = mhc_tilelang.current_platform.device_type + n_splits = 8 + gemm_mul = torch.randn((n_splits, 1, 24), device=device, dtype=torch.float32).mul_( + 0.01 + ) + gemm_sqrsum = torch.rand((n_splits, 1), device=device, dtype=torch.float32).add_( + 4096 + ) + scale = torch.tensor([0.5, 0.75, 0.25], device=device, dtype=torch.float32) + base = torch.randn((24,), device=device, dtype=torch.float32).mul_(0.1) + residual = torch.randn((1, 4, 4096), device=device, dtype=torch.float16).mul_(0.1) + norm_weight = torch.randn((4096,), device=device, dtype=torch.float16).mul_(0.1) + + def run( + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + layer_input: torch.Tensor, + ) -> None: + sm70_mhc_pre_norm_from_staging( + gemm_mul, + gemm_sqrsum, + scale, + base, + residual, + post_mix, + comb_mix, + layer_input, + norm_weight, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 20, + 1e-6, + ) + + eager_outputs = ( + torch.empty((1, 4), device=device, dtype=torch.float32), + torch.empty((1, 4, 4), device=device, dtype=torch.float32), + torch.empty((1, 4096), device=device, dtype=torch.float16), + ) + run(*eager_outputs) + expected = tuple(output.clone() for output in eager_outputs) + + graph_outputs = tuple(torch.empty_like(output) for output in eager_outputs) + run(*graph_outputs) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(*graph_outputs) + graph.replay() + torch.accelerator.synchronize() + + for output, reference in zip(graph_outputs, expected): + torch.testing.assert_close(output, reference, rtol=0, atol=0) + + +@pytest.mark.parametrize("num_tokens", [1, 4, 16, 17]) +@pytest.mark.skipif( + not ( + mhc_tilelang.current_platform.is_cuda() + and has_tilelang() + and mhc_tilelang.current_platform.get_device_capability() + == DeviceCapability(7, 0) + ), + reason="NVIDIA V100/SM70 and TileLang required", +) +def test_mhc_sm70_fused_fast_paths_match_torch(num_tokens: int) -> None: + torch.manual_seed(20260827 + num_tokens) + hidden_size = 4096 + hc_mult = 4 + device = mhc_tilelang.current_platform.device_type + + x = torch.randn((num_tokens, hidden_size), device=device, dtype=torch.float16).mul_( + 0.1 + ) + residual = torch.randn( + (num_tokens, hc_mult, hidden_size), device=device, dtype=torch.float16 + ).mul_(0.1) + post_mix = torch.sigmoid( + torch.randn((num_tokens, hc_mult, 1), device=device, dtype=torch.float32) + ) + comb_mix = torch.softmax( + torch.randn((num_tokens, hc_mult, hc_mult), device=device, dtype=torch.float32), + dim=-1, + ) + fn = torch.randn( + (24, hc_mult * hidden_size), device=device, dtype=torch.float32 + ).mul_(0.01) + scale = torch.tensor([0.5, 0.75, 0.25], device=device, dtype=torch.float32) + base = torch.randn((24,), device=device, dtype=torch.float32).mul_(0.1) + norm_weight = torch.randn((hidden_size,), device=device, dtype=torch.float16).mul_( + 0.1 + ) + + actual = mhc_tilelang.mhc_fused_post_pre_tilelang( + x, + residual, + post_mix, + comb_mix, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 20, + norm_weight=norm_weight, + norm_eps=1e-6, + ) + expected = mhc_fused_post_pre_torch( + x, + residual, + post_mix, + comb_mix, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 20, + norm_weight=norm_weight, + norm_eps=1e-6, + ) + + tolerances = ((2e-3, 2e-3), (2e-3, 2e-4), (3e-3, 3e-4), (5e-3, 5e-4)) + for output, reference, (rtol, atol) in zip(actual, expected, tolerances): + torch.testing.assert_close(output, reference, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("num_tokens", [1, 17]) +@pytest.mark.skipif( + not ( + mhc_tilelang.current_platform.is_cuda() + and has_tilelang() + and mhc_tilelang.current_platform.get_device_capability() + == DeviceCapability(7, 0) + ), + reason="NVIDIA V100/SM70 and TileLang required", +) +def test_mhc_sm70_standalone_pre_post_fast_paths_match_torch( + num_tokens: int, +) -> None: + torch.manual_seed(20260828 + num_tokens) + hidden_size = 4096 + hc_mult = 4 + device = mhc_tilelang.current_platform.device_type + x = torch.randn((num_tokens, hidden_size), device=device, dtype=torch.float16).mul_( + 0.1 + ) + residual = torch.randn( + (num_tokens, hc_mult, hidden_size), device=device, dtype=torch.float16 + ).mul_(0.1) + post_mix = torch.sigmoid( + torch.randn((num_tokens, hc_mult, 1), device=device, dtype=torch.float32) + ) + comb_mix = torch.softmax( + torch.randn((num_tokens, hc_mult, hc_mult), device=device, dtype=torch.float32), + dim=-1, + ) + fn = torch.randn( + (24, hc_mult * hidden_size), device=device, dtype=torch.float32 + ).mul_(0.01) + scale = torch.tensor([0.5, 0.75, 0.25], device=device, dtype=torch.float32) + base = torch.randn((24,), device=device, dtype=torch.float32).mul_(0.1) + norm_weight = torch.randn((hidden_size,), device=device, dtype=torch.float16).mul_( + 0.1 + ) + + actual_post = mhc_tilelang.mhc_post_tilelang(x, residual, post_mix, comb_mix) + expected_post = mhc_post_torch(x, residual, post_mix, comb_mix) + actual_pre = mhc_tilelang.mhc_pre_tilelang( + residual, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 20, + norm_weight=norm_weight, + norm_eps=1e-6, + ) + expected_pre = mhc_pre_torch( + residual, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 20, + norm_weight=norm_weight, + norm_eps=1e-6, + ) + + torch.testing.assert_close(actual_post, expected_post, rtol=2e-3, atol=2e-4) + tolerances = ((2e-3, 2e-4), (3e-3, 3e-4), (5e-3, 1e-3)) + for output, reference, (rtol, atol) in zip(actual_pre, expected_pre, tolerances): + torch.testing.assert_close(output, reference, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("num_tokens", [1, 17]) +@pytest.mark.skipif( + not ( + mhc_tilelang.current_platform.is_cuda() + and has_tilelang() + and mhc_tilelang.current_platform.get_device_capability() + == DeviceCapability(7, 0) + ), + reason="NVIDIA V100/SM70 and TileLang required", +) +def test_mhc_sm70_broadcast_pre_fast_path_matches_expanded_torch( + num_tokens: int, +) -> None: + torch.manual_seed(20260830 + num_tokens) + hidden_size = 4096 + hc_mult = 4 + device = mhc_tilelang.current_platform.device_type + x = torch.randn((num_tokens, hidden_size), device=device, dtype=torch.float16).mul_( + 0.1 + ) + residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() + fn = torch.randn( + (24, hc_mult * hidden_size), device=device, dtype=torch.float32 + ).mul_(0.01) + fn_broadcast = fn.view(24, hc_mult, hidden_size).sum(dim=1) + scale = torch.tensor([0.5, 0.75, 0.25], device=device, dtype=torch.float32) + base = torch.randn((24,), device=device, dtype=torch.float32).mul_(0.1) + norm_weight = torch.randn((hidden_size,), device=device, dtype=torch.float16).mul_( + 0.1 + ) + + actual = mhc_tilelang.mhc_pre_broadcast_tilelang( + x, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 20, + norm_weight=norm_weight, + norm_eps=1e-6, + fn_broadcast=fn_broadcast, + ) + expected = mhc_pre_torch( + residual, + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 20, + norm_weight=norm_weight, + norm_eps=1e-6, + ) + + torch.testing.assert_close(actual[0], residual, rtol=0, atol=0) + tolerances = ((2e-3, 2e-4), (3e-3, 3e-4), (5e-3, 1e-3)) + for output, reference, (rtol, atol) in zip(actual[1:], expected, tolerances): + torch.testing.assert_close(output, reference, rtol=rtol, atol=atol) + + @pytest.mark.skipif( - not (mhc_tilelang.current_platform.is_cuda() and has_tilelang()), - reason="CUDA and TileLang required", + not (mhc_tilelang.current_platform.is_cuda() and has_tilelang()) + or mhc_tilelang.current_platform.get_device_capability() == DeviceCapability(7, 0), + reason="SM70 uses the FP16 torch fallback; native TileLang CUDA required", ) -def test_mhc_sm70_fp16_block_m_prenorm_keeps_rows_independent() -> None: +def test_mhc_fp16_tilelang_block_m_prenorm_keeps_rows_independent() -> None: """Exercise the >=1024-token block-M path with non-identical row pairs.""" hidden_size = 128 hc_mult = 4 diff --git a/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py b/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py index b5297317c0..f3ed2c5669 100644 --- a/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py +++ b/tests/kernels/test_sm70_deepseek_v4_fp16_gemv.py @@ -8,6 +8,9 @@ import vllm.envs as envs from vllm.models.deepseek_v4.sm70.gemv import ( + _has_sm70_dsv4_fp13_weight_contract, + _pack_sm70_dsv4_fp13_weight, + can_use_sm70_dsv4_fp13_gemv, can_use_sm70_dsv4_fp16_gemv, maybe_sm70_dsv4_fp16_gemv, ) @@ -20,13 +23,79 @@ def reset_env_cache() -> Iterator[None]: envs.disable_envs_cache() -def test_sm70_dsv4_fp16_gemv_is_default_off(monkeypatch) -> None: +def test_sm70_dsv4_fp13_gemv_is_default_on_with_rollback(monkeypatch) -> None: monkeypatch.delenv("VLLM_SM70_DSV4_FP16_GEMV", raising=False) + monkeypatch.delenv("VLLM_SM70_DSV4_FP13_GEMV", raising=False) + assert not envs.VLLM_SM70_DSV4_FP16_GEMV + assert envs.VLLM_SM70_DSV4_FP13_GEMV + + monkeypatch.setenv("VLLM_SM70_DSV4_FP13_GEMV", "0") + envs.disable_envs_cache() + assert not envs.VLLM_SM70_DSV4_FP13_GEMV + + +def test_sm70_dsv4_fp13_enables_fp16_fallback(monkeypatch) -> None: + monkeypatch.setenv("VLLM_SM70_DSV4_FP16_GEMV", "0") + monkeypatch.setenv("VLLM_SM70_DSV4_FP13_GEMV", "1") + monkeypatch.setattr( + "vllm.models.deepseek_v4.sm70.gemv._has_sm70_dsv4_gemv_contract", + lambda *args: True, + ) + x = torch.empty((1, 4096), dtype=torch.float16) + weight = torch.empty((256, 4096), dtype=torch.float16) + assert can_use_sm70_dsv4_fp16_gemv(x, weight, torch.float32) + + +def test_sm70_dsv4_gemv_rejects_cpu_tensors(monkeypatch) -> None: + monkeypatch.setenv("VLLM_SM70_DSV4_FP13_GEMV", "1") + monkeypatch.setattr( + "vllm.models.deepseek_v4.sm70.gemv.current_platform.is_cuda", + lambda: True, + ) + monkeypatch.setattr( + "vllm.models.deepseek_v4.sm70.gemv.current_platform.is_device_capability", + lambda capability: capability == (7, 0), + ) x = torch.empty((1, 4096), dtype=torch.float16) weight = torch.empty((256, 4096), dtype=torch.float16) assert not can_use_sm70_dsv4_fp16_gemv(x, weight, torch.float32) +def test_sm70_dsv4_fp13_weight_contract_is_tensor_based() -> None: + compatible_bits = torch.tensor([0x0000, 0x0007, 0x3C00, -0x4400], dtype=torch.int16) + assert _has_sm70_dsv4_fp13_weight_contract(compatible_bits.view(torch.float16)) + + normal_with_discarded_bit = torch.tensor([0x3C01], dtype=torch.int16) + assert not _has_sm70_dsv4_fp13_weight_contract( + normal_with_discarded_bit.view(torch.float16) + ) + nonfinite = torch.tensor([0x7C00], dtype=torch.int16) + assert not _has_sm70_dsv4_fp13_weight_contract(nonfinite.view(torch.float16)) + + +def test_sm70_dsv4_fp13_pack_preserves_upper_13_bits() -> None: + raw = (torch.arange(64 * 4096, dtype=torch.int32).mul(40503).add(17) & 0xFFFF).to( + torch.uint16 + ) + weight = raw.view(torch.float16).reshape(64, 4096) + packed = _pack_sm70_dsv4_fp13_weight(weight) + assert packed.shape == (64, 128, 13) + + words = packed.to(torch.int64) & 0xFFFFFFFF + unpacked = [] + for value_index in range(32): + bit_offset = value_index * 13 + word_index = bit_offset // 32 + shift = bit_offset % 32 + code = words[..., word_index] >> shift + if shift > 19: + code |= words[..., word_index + 1] << (32 - shift) + unpacked.append(code & 0x1FFF) + codes = torch.stack(unpacked, dim=-1).reshape(64, 4096) + expected = (raw.to(torch.int32) >> 3).reshape(64, 4096) + assert torch.equal(codes.to(torch.int32), expected) + + @pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), reason="requires NVIDIA V100/SM70", @@ -70,3 +139,54 @@ def test_sm70_dsv4_fp16_gemv_graph( else: reference = torch.mm(x, weight.T, out_dtype=torch.float32) torch.testing.assert_close(captured, reference, rtol=2e-4, atol=5e-5) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), + reason="requires NVIDIA V100/SM70", +) +@pytest.mark.parametrize( + ("n", "output_dtype"), + [ + (64, torch.float16), + (256, torch.float32), + (512, torch.float32), + (1024, torch.float32), + (2048, torch.float32), + ], +) +def test_sm70_dsv4_fp13_gemv_graph( + monkeypatch, n: int, output_dtype: torch.dtype +) -> None: + monkeypatch.setenv("VLLM_SM70_DSV4_FP16_GEMV", "0") + monkeypatch.setenv("VLLM_SM70_DSV4_FP13_GEMV", "1") + torch.manual_seed(20260826 + n) + x = torch.randn((1, 4096), device="cuda", dtype=torch.float16) + source = torch.randn((n, 4096), device="cuda", dtype=torch.bfloat16) * 0.01 + weight = source.to(torch.float16) + packed = _pack_sm70_dsv4_fp13_weight(weight) + assert can_use_sm70_dsv4_fp13_gemv(x, weight, packed, output_dtype) + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + candidate = maybe_sm70_dsv4_fp16_gemv( + x, weight, output_dtype, packed_weight=packed + ) + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + assert candidate is not None + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = maybe_sm70_dsv4_fp16_gemv( + x, weight, output_dtype, packed_weight=packed + ) + assert captured is not None + + x.copy_(torch.randn_like(x)) + graph.replay() + torch.cuda.synchronize() + reference = torch.mm(x, weight.T, out_dtype=torch.float32) + if output_dtype == torch.float16: + reference = reference.to(torch.float16) + torch.testing.assert_close(captured, reference, rtol=3e-4, atol=5e-5) diff --git a/tests/model_executor/test_qwen3_5_quantization.py b/tests/model_executor/test_qwen3_5_quantization.py index 74914bd765..b5b46422f5 100644 --- a/tests/model_executor/test_qwen3_5_quantization.py +++ b/tests/model_executor/test_qwen3_5_quantization.py @@ -70,18 +70,19 @@ def test_sm70_gdn_qpn8_ba_split_accepts_complete_contract(monkeypatch) -> None: assert gdn._sm70_gdn_qpn8_ba_split_enabled() -def test_sm70_gdn_qpn8_ba_weight_contract_is_layout_based() -> None: +@pytest.mark.parametrize("input_width", [2560, 5120]) +def test_sm70_gdn_qpn8_ba_weight_contract_is_layout_based(input_width) -> None: from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gdn qkvz = SimpleNamespace( sm70_fp8_qpn8=True, sm70_fp8_prefill_exact_dense_workspace_ptr=42, - weight=torch.empty((5120, 4096), dtype=torch.uint8, device="meta"), + weight=torch.empty((input_width, 4096), dtype=torch.uint8, device="meta"), weight_scale_inv=torch.empty((1, 4096), dtype=torch.float16, device="meta"), bias=None, ) ba = SimpleNamespace( - weight=torch.empty((24, 5120), dtype=torch.float16, device="meta"), + weight=torch.empty((24, input_width), dtype=torch.float16, device="meta"), bias=None, ) layer = SimpleNamespace( @@ -95,7 +96,7 @@ def test_sm70_gdn_qpn8_ba_weight_contract_is_layout_based() -> None: qkvz.bias = object() assert not gdn._sm70_gdn_qpn8_ba_weight_contract(layer) qkvz.bias = None - qkvz.weight = torch.empty((4096, 5120), dtype=torch.uint8, device="meta") + qkvz.weight = torch.empty((4096, input_width), dtype=torch.uint8, device="meta") assert not gdn._sm70_gdn_qpn8_ba_weight_contract(layer) diff --git a/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py index a31aba81ae..154ec2626e 100644 --- a/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py +++ b/tests/model_executor/test_sm70_fp8_qpn8_pp2_tp4.py @@ -30,11 +30,15 @@ def _layer( ) -def test_pp2_tp4_qpn8_is_default_on_with_explicit_rollback(monkeypatch) -> None: +def test_pp2_tp4_qpn8_is_default_off_with_explicit_opt_in(monkeypatch) -> None: monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4", raising=False) envs.disable_envs_cache() try: + assert not fp8._sm70_fp8_qpn8_pp2_tp4_enabled() + + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "1") + envs.disable_envs_cache() assert fp8._sm70_fp8_qpn8_pp2_tp4_enabled() monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "0") @@ -49,13 +53,31 @@ def test_pp2_tp4_qpn8_is_default_on_with_explicit_rollback(monkeypatch) -> None: envs.disable_envs_cache() +def test_qpn8_extension_load_requires_explicit_route(monkeypatch) -> None: + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_LIBRARY", "/tmp/qpn8-test.so") + monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) + monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4", raising=False) + with patch.object(torch.ops, "load_library") as load_library: + fp8.sm70_ops._maybe_load_fp8_qpn8_library() + load_library.assert_not_called() + + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "1") + fp8.sm70_ops._maybe_load_fp8_qpn8_library() + load_library.assert_called_once_with("/tmp/qpn8-test.so") + + load_library.reset_mock() + monkeypatch.setenv("VLLM_SM70_FP8_QPN8", "0") + fp8.sm70_ops._maybe_load_fp8_qpn8_library() + load_library.assert_not_called() + + @pytest.mark.parametrize( ("layer", "gated_silu", "expected"), [ - (_layer("fused_wqa_wkv", 1, 4096, 1536), False, (16, 1, True)), - (_layer("wq_b", 4, 1024, 8192), False, (4, 1, False)), - (_layer("wo_b", 4, 2048, 4096), False, (8, 2, False)), - (_layer("down_proj", 4, 512, 4096), False, (4, 1, False)), + (_layer("fused_wqa_wkv", 1, 4096, 1536), False, (32, 2, False)), + (_layer("wq_b", 4, 1024, 8192), False, (8, 2, False)), + (_layer("wo_b", 4, 2048, 4096), False, (16, 2, False)), + (_layer("down_proj", 4, 512, 4096), False, (16, 2, False)), ], ) def test_pp2_tp4_qpn8_exact_operator_contracts( @@ -90,6 +112,43 @@ def test_pp2_tp4_qpn8_rejects_wrong_tensor_and_concurrency_roles() -> None: assert fp8._sm70_fp8_qpn8_pp2_tp4_config(wrong_layout, gated_silu=False) is None +def test_pp2_tp4_qpn8_shared_gate_requires_explicit_opt_in(monkeypatch) -> None: + layer = _layer( + "gate_up_proj", + 4, + 4096, + 1024, + output_partition_sizes=[512, 512], + ) + layer.prefix = "arbitrary.layers.7.mlp.shared_experts.gate_up_proj" + + monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", raising=False) + envs.disable_envs_cache() + try: + assert fp8._is_sm70_fp8_qpn8_pp2_tp4_shared_gate_contract(layer) + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=False) is None + + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", "1") + envs.disable_envs_cache() + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=False) == ( + 32, + 2, + False, + ) + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=True) is None + + layer.prefix = "arbitrary.layers.7.mlp.gate_up_proj" + assert not fp8._is_sm70_fp8_qpn8_pp2_tp4_shared_gate_contract(layer) + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=False) is None + + layer.prefix = "arbitrary.layers.7.mlp.shared_experts.gate_up_proj" + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", "0") + envs.disable_envs_cache() + assert fp8._sm70_fp8_qpn8_pp2_tp4_config(layer, gated_silu=False) is None + finally: + envs.disable_envs_cache() + + def test_pp2_tp4_qpn8_runtime_and_workspace_contract() -> None: config = SimpleNamespace( parallel_config=SimpleNamespace( @@ -125,16 +184,16 @@ def test_pp2_tp4_qpn8_grouped_dispatches_caller_groups() -> None: layer = _layer("wo_a", 4, 4096, 2048) layer.is_bmm = True layer.bmm_batch_size = 2 - assert fp8._sm70_fp8_qpn8_pp2_tp4_bmm_config(layer) == (16, 1, True) + assert fp8._sm70_fp8_qpn8_pp2_tp4_bmm_config(layer) == (32, 2, False) layer.sm70_fp8_turbomind = True layer.sm70_fp8_qpn8 = True layer.sm70_fp8_qpn8_bmm = True layer.sm70_fp8_bmm_groups = 2 layer.sm70_fp8_bmm_output_size = 1024 - layer.sm70_fp8_qpn8_split_k = 16 - layer.sm70_fp8_qpn8_nacc = 1 - layer.sm70_fp8_qpn8_prefetch = True + layer.sm70_fp8_qpn8_split_k = 32 + layer.sm70_fp8_qpn8_nacc = 2 + layer.sm70_fp8_qpn8_prefetch = False layer.sm70_fp8_prefill_exact_dense_workspace_ptr = 123 layer.weight = torch.empty((2, 4096, 1024), device="meta") layer.weight_scale_inv = torch.empty((2, 256, 32), device="meta") @@ -157,24 +216,24 @@ def fake_dispatch( (tuple(input_.shape), workspace_ptr, split_k, prefetch, gated_silu) ) out.fill_(len(calls)) - assert nacc == 1 + assert nacc == 2 x = torch.zeros((1, 2, 4096), dtype=torch.float16) with patch.object(fp8.sm70_ops, "fp8_qpn8_dispatch_sm70_out", fake_dispatch): out = fp8.Fp8LinearMethod.apply(None, layer, x) assert calls == [ - ((1, 4096), 123, 16, True, False), - ((1, 4096), 123, 16, True, False), + ((1, 4096), 123, 32, False, False), + ((1, 4096), 123, 32, False, False), ] assert out.shape == (1, 2, 1024) torch.testing.assert_close(out[:, 0], torch.ones_like(out[:, 0])) torch.testing.assert_close(out[:, 1], torch.full_like(out[:, 1], 2)) -def test_pp2_tp4_qpn8_default_prepares_matching_layer(monkeypatch) -> None: +def test_pp2_tp4_qpn8_explicit_opt_in_prepares_matching_layer(monkeypatch) -> None: monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) - monkeypatch.delenv("VLLM_SM70_FP8_QPN8_PP2_TP4", raising=False) + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "1") envs.disable_envs_cache() layer = _layer("fused_wqa_wkv", 1, 4096, 1536) layer.orig_dtype = torch.float16 @@ -224,6 +283,72 @@ def test_pp2_tp4_qpn8_default_prepares_matching_layer(monkeypatch) -> None: envs.disable_envs_cache() assert layer.sm70_fp8_qpn8 - assert layer.sm70_fp8_qpn8_split_k == 16 - assert layer.sm70_fp8_qpn8_nacc == 1 - assert layer.sm70_fp8_qpn8_prefetch + assert layer.sm70_fp8_qpn8_split_k == 32 + assert layer.sm70_fp8_qpn8_nacc == 2 + assert not layer.sm70_fp8_qpn8_prefetch + + +def test_pp2_tp4_qpn8_shared_gate_retains_external_activation(monkeypatch) -> None: + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "1") + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", "1") + monkeypatch.delenv("VLLM_SM70_FP8_QPN8", raising=False) + envs.disable_envs_cache() + layer = _layer( + "gate_up_proj", + 4, + 4096, + 1024, + output_partition_sizes=[512, 512], + ) + layer.prefix = "arbitrary.layers.7.mlp.shared_experts.gate_up_proj" + layer.orig_dtype = torch.float16 + layer.is_bmm = False + layer.weight_scale_inv = torch.empty((8, 32), device="meta") + method = fp8.Fp8LinearMethod.__new__(fp8.Fp8LinearMethod) + method.use_marlin = False + method.use_sm70_fp8_turbomind = True + method.weight_block_size = [128, 128] + config = SimpleNamespace( + parallel_config=SimpleNamespace( + pipeline_parallel_size=2, + tensor_parallel_size=4, + enable_dbo=False, + ubatch_size=0, + ), + scheduler_config=SimpleNamespace(max_num_seqs=1), + speculative_config=None, + ) + + try: + with ( + patch.object(fp8, "get_current_vllm_config", return_value=config), + patch.object(fp8, "_missing_sm70_fp8_qpn8_ops", return_value=[]), + patch.object( + fp8, + "_get_sm70_fp8_qpn8_pp2_tp4_workspace", + return_value=torch.empty(1), + ), + patch.object( + fp8, + "process_fp8_weight_block_strategy", + side_effect=lambda weight, scales: (weight, scales), + ), + patch.object( + fp8.sm70_ops, + "fp8_qpn8_prepare_sm70", + return_value=( + torch.empty((4096, 1024), dtype=torch.uint8, device="meta"), + torch.empty((256, 32), dtype=torch.float16, device="meta"), + ), + ), + patch.object(fp8, "replace_parameter"), + ): + method.process_weights_after_loading(layer) + finally: + envs.disable_envs_cache() + + assert layer.sm70_fp8_qpn8 + assert layer.sm70_fp8_qpn8_split_k == 32 + assert layer.sm70_fp8_qpn8_nacc == 2 + assert not layer.sm70_fp8_qpn8_prefetch + assert not getattr(layer, "sm70_fp8_gated_silu", False) diff --git a/tests/model_executor/test_weights_mapper_stacked.py b/tests/model_executor/test_weights_mapper_stacked.py new file mode 100644 index 0000000000..65b655c730 --- /dev/null +++ b/tests/model_executor/test_weights_mapper_stacked.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, +) +from vllm.model_executor.models.utils import WeightsMapper + + +def test_weights_mapper_preserves_stacked_shard_id() -> None: + mapper = WeightsMapper( + orig_to_new_stacked={ + ".input_mix_weight_down.weight": ( + ".input_mix_weight_down_block_inject.weight", + 0, + ), + ".block_inject_weight.weight": ( + ".input_mix_weight_down_block_inject.weight", + 1, + ), + } + ) + down = torch.ones(2, 4) + injection = torch.full((3, 4), 2.0) + + mapped = list( + mapper.apply( + [ + ("layer.input_mix_weight_down.weight", down), + ("layer.block_inject_weight.weight", injection), + ] + ) + ) + + assert [name for name, _ in mapped] == [ + "layer.input_mix_weight_down_block_inject.weight", + "layer.input_mix_weight_down_block_inject.weight", + ] + assert down.shard_id == 0 + assert injection.shard_id == 1 + + +def test_weights_mapper_preserves_named_stacked_shard_id() -> None: + mapper = WeightsMapper( + orig_to_new_stacked={ + ".q.": (".qkv.", "q"), + ".k.": (".qkv.", "k"), + } + ) + q = torch.ones(2, 4) + k = torch.full((2, 4), 2.0) + + mapped = list( + mapper.apply( + [ + ("layer.q.weight", q), + ("layer.k.weight", k), + ] + ) + ) + + assert [name for name, _ in mapped] == [ + "layer.qkv.weight", + "layer.qkv.weight", + ] + assert q.shard_id == "q" + assert k.shard_id == "k" + + +def test_merged_column_load_weights_forwards_stacked_shards() -> None: + layer = object.__new__(MergedColumnParallelLinear) + torch.nn.Module.__init__(layer) + layer.output_sizes = [2, 3] + layer.tp_size = 1 + layer.tp_rank = 0 + layer.prefix = "hc.input_mix_weight_down_block_inject" + weight = torch.nn.Parameter(torch.zeros(5, 4)) + calls = [] + + def weight_loader(param, loaded_weight, shard_id) -> None: + calls.append((param, loaded_weight, shard_id)) + + weight.weight_loader = weight_loader + layer.register_parameter("weight", weight) + down = torch.ones(2, 4) + down.shard_id = 0 + injection = torch.full((3, 4), 2.0) + injection.shard_id = 1 + + loaded = list( + layer.load_weights( + [ + ("weight", down), + ("weight", injection), + ] + ) + ) + + assert loaded == ["weight", "weight"] + assert calls == [ + (weight, down, 0), + (weight, injection, 1), + ] + + +def test_qkv_load_weights_forwards_stacked_shards() -> None: + layer = object.__new__(QKVParallelLinear) + torch.nn.Module.__init__(layer) + layer.prefix = "attn.qkv_proj" + weight = torch.nn.Parameter(torch.zeros(7, 4)) + calls = [] + + def weight_loader(param, loaded_weight, shard_id) -> None: + calls.append((param, loaded_weight, shard_id)) + + weight.weight_loader = weight_loader + layer.register_parameter("weight", weight) + query = torch.ones(3, 4) + query.shard_id = "q" + key = torch.full((2, 4), 2.0) + key.shard_id = "k" + value = torch.full((2, 4), 3.0) + value.shard_id = "v" + + loaded = list( + layer.load_weights( + [ + ("weight", query), + ("weight", key), + ("weight", value), + ] + ) + ) + + assert loaded == ["weight", "weight", "weight"] + assert calls == [ + (weight, query, "q"), + (weight, key, "k"), + (weight, value, "v"), + ] diff --git a/tests/models/glm5next/test_pp.py b/tests/models/glm5next/test_pp.py new file mode 100644 index 0000000000..b7d1346d69 --- /dev/null +++ b/tests/models/glm5next/test_pp.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch +from torch import nn + +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.models.glm5next.nvidia.kda import Glm5NextLinearAttention +from vllm.models.glm5next.nvidia.model import ( + Glm5NextForConditionalGeneration, + Glm5NextModel, +) + + +def test_glm53_pp_intermediate_state_shapes(): + model = Glm5NextModel.__new__(Glm5NextModel) + nn.Module.__init__(model) + model.config = SimpleNamespace( + hidden_size=4096, mhc=True, mhc_num_residual_streams=4 + ) + + tensors = model.make_empty_intermediate_tensors( + batch_size=7, dtype=torch.float16, device=torch.device("cpu") + ).tensors + assert tuple(tensors) == ("hidden_states", "residual", "post", "comb") + assert tensors["hidden_states"].shape == (7, 4096) + assert tensors["residual"].shape == (7, 4, 4096) + assert tensors["post"].shape == (7, 4, 1) + assert tensors["comb"].shape == (7, 4, 4) + assert tensors["hidden_states"].dtype == torch.float16 + assert tensors["residual"].dtype == torch.float16 + assert tensors["post"].dtype == torch.float32 + assert tensors["comb"].dtype == torch.float32 + + +def test_glm53_merged_kda_state_contract(): + vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=4), + speculative_config=None, + model_config=SimpleNamespace( + dtype=torch.float16, + hf_config=SimpleNamespace( + linear_num_heads=64, + linear_head_dim=128, + linear_conv_kernel_dim=4, + ), + ), + cache_config=SimpleNamespace(mamba_cache_dtype="auto"), + ) + model_shapes = Glm5NextForConditionalGeneration.get_mamba_state_shape_from_config( + vllm_config + ) + model_dtypes = Glm5NextForConditionalGeneration.get_mamba_state_dtype_from_config( + vllm_config + ) + copy_funcs = Glm5NextForConditionalGeneration.get_mamba_state_copy_func() + + layer = Glm5NextLinearAttention.__new__(Glm5NextLinearAttention) + nn.Module.__init__(layer) + layer.tp_size = 4 + layer.num_heads = 64 + layer.head_dim = 128 + layer.conv_size = 4 + layer.num_spec = 0 + layer.model_config = vllm_config.model_config + layer.cache_config = vllm_config.cache_config + + assert layer.get_state_shape() == model_shapes + assert layer.get_state_dtype() == model_dtypes + assert model_dtypes == (torch.float16, torch.float32) + assert len(copy_funcs) == 2 + assert model_shapes[1] == (16, 128, 128) + expected_conv = (6144, 3) if is_conv_state_dim_first() else (3, 6144) + assert model_shapes[0] == expected_conv diff --git a/tests/models/glm5next/test_sm70_kda.py b/tests/models/glm5next/test_sm70_kda.py new file mode 100644 index 0000000000..76d155e1d7 --- /dev/null +++ b/tests/models/glm5next/test_sm70_kda.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +import vllm._sm70_ops as sm70_ops +from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability + + +@pytest.mark.skipif( + not ( + current_platform.is_cuda() + and current_platform.get_device_capability() == DeviceCapability(7, 0) + and hasattr(torch.ops._C, "sm70_glm_kda_fg_b_out") + ), + reason="native NVIDIA V100/SM70 GLM KDA CUDA op required", +) +def test_sm70_glm_kda_fused_fg_b_matches_fp16_and_graph() -> None: + torch.manual_seed(20260827) + device = current_platform.device_type + f_input = torch.randn((1, 128), device=device, dtype=torch.float16).mul_(0.1) + g_input = torch.randn((1, 128), device=device, dtype=torch.float16).mul_(0.1) + f_weight = torch.randn((2048, 128), device=device, dtype=torch.float16).mul_(0.01) + g_weight = torch.randn((2048, 128), device=device, dtype=torch.float16).mul_(0.01) + f_out = torch.empty((1, 2048), device=device, dtype=torch.float16) + g_out = torch.empty_like(f_out) + + def run() -> None: + sm70_ops.sm70_glm_kda_fg_b_out( + f_out, g_out, f_input, g_input, f_weight, g_weight + ) + + run() + torch.accelerator.synchronize() + expected = (F.linear(f_input, f_weight), F.linear(g_input, g_weight)) + torch.testing.assert_close(f_out, expected[0], rtol=2e-3, atol=2e-4) + torch.testing.assert_close(g_out, expected[1], rtol=2e-3, atol=2e-4) + eager = (f_out.clone(), g_out.clone()) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + graph.replay() + torch.accelerator.synchronize() + + torch.testing.assert_close(f_out, eager[0], rtol=0, atol=0) + torch.testing.assert_close(g_out, eager[1], rtol=0, atol=0) diff --git a/tests/models/glm5next/test_sm70_sparse.py b/tests/models/glm5next/test_sm70_sparse.py new file mode 100644 index 0000000000..4427e2ad3b --- /dev/null +++ b/tests/models/glm5next/test_sm70_sparse.py @@ -0,0 +1,336 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.attention.mla_attention import MLAAttention +from vllm.models.common.ops import fused_q_kv_rmsnorm +from vllm.models.deepseek_v4.sm70.sparse_kernels import ( + sm70_sparse_attention_gathered, +) +from vllm.models.glm5next.nvidia.ops.kpool_compress import ( + _hadamard128_torch, + fwht128, +) +from vllm.models.glm5next.sm70 import sparse as sparse_module +from vllm.models.glm5next.sm70.fp8_kv import ( + sm70_glm5_fp8_kv_insert, + sm70_glm5_sparse_attention_paged_fp8, + sm70_glm5_sparse_attention_paged_fp8_gemm, +) +from vllm.models.glm5next.sm70.sparse import Glm5NextSM70SparseBackend +from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability +from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig +from vllm.transformers_utils.model_arch_config_convertor import ( + ModelArchConfigConvertorBase, +) +from vllm.v1.kv_cache_interface import MLAAttentionSpec + + +def test_glm53_model_arch_uses_mla_cache_width(): + config = Glm5NextConfig( + architectures=["Glm5NextForConditionalGeneration"], + text_config={ + "head_dim": 0, + "kv_lora_rank": 512, + "qk_rope_head_dim": 0, + }, + ) + converted = ModelArchConfigConvertorBase(config, config.text_config).convert() + + assert converted.is_deepseek_mla + assert converted.head_size == 512 + + +def test_glm53_sm70_sparse_backend_contract(): + invalid = Glm5NextSM70SparseBackend.validate_configuration( + head_size=512, + dtype=torch.float16, + kv_cache_dtype="auto", + block_size=1152, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(major=7, minor=0), + attn_type="decoder", + ) + assert invalid == [] + + +def test_glm53_sm70_sparse_backend_uses_structural_gates(monkeypatch): + text_config = SimpleNamespace( + model_type="generic_sparse_mla", + qk_rope_head_dim=0, + kv_lora_rank=512, + ) + monkeypatch.setattr( + sparse_module, + "get_current_vllm_config_or_none", + lambda: SimpleNamespace( + model_config=SimpleNamespace(hf_text_config=text_config) + ), + ) + + assert ( + Glm5NextSM70SparseBackend.supports_combination( + head_size=512, + dtype=torch.float16, + kv_cache_dtype="auto", + block_size=1152, + use_mla=True, + has_sink=False, + use_sparse=True, + device_capability=DeviceCapability(major=7, minor=0), + ) + is None + ) + + text_config.qk_rope_head_dim = 64 + assert "requires NoPE" in ( + Glm5NextSM70SparseBackend.supports_combination( + head_size=512, + dtype=torch.float16, + kv_cache_dtype="auto", + block_size=1152, + use_mla=True, + has_sink=False, + use_sparse=True, + device_capability=DeviceCapability(major=7, minor=0), + ) + or "" + ) + + +def test_glm53_sm70_fp8_kv_layout_contract(): + invalid = Glm5NextSM70SparseBackend.validate_configuration( + head_size=512, + dtype=torch.float16, + kv_cache_dtype="fp8_e4m3", + block_size=1152, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(major=7, minor=0), + attn_type="decoder", + ) + assert invalid == [] + assert Glm5NextSM70SparseBackend.get_kv_cache_shape( + num_blocks=3, + block_size=64, + num_kv_heads=1, + head_size=512, + cache_dtype_str="fp8_e4m3", + ) == (3, 64, 520) + + spec = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=512, + dtype=torch.uint8, + cache_dtype_str="fp8_e4m3", + model_version="glm5_next", + ) + assert spec.real_page_size_bytes == 64 * 520 + + layer = MLAAttention.__new__(MLAAttention) + torch.nn.Module.__init__(layer) + layer.kv_cache_dtype = "fp8_e4m3" + layer.head_size = 512 + layer.attn_backend = Glm5NextSM70SparseBackend + emitted = layer.get_kv_cache_spec( + SimpleNamespace( + cache_config=SimpleNamespace(block_size=64, cache_dtype="fp8_e4m3"), + model_config=SimpleNamespace(dtype=torch.float16), + ) + ) + assert isinstance(emitted, MLAAttentionSpec) + assert emitted.model_version == "glm5_next" + assert emitted.real_page_size_bytes == 64 * 520 + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.is_device_capability((7, 0)), + reason="requires an exact SM70 CUDA device", +) +def test_glm53_sm70_sparse_attention_matches_reference(): + torch.manual_seed(7) + device = torch.device("cuda") + q = torch.randn(2, 16, 512, dtype=torch.float16, device=device) * 0.1 + kv = torch.randn(48, 512, dtype=torch.float16, device=device) * 0.1 + indices = torch.tensor( + [ + [0, 3, 5, 9, 11, 17, 23, 31] + [-1] * 8, + [2, 4, 8, 12, 18, 24, 32, 40, 45] + [-1] * 7, + ], + dtype=torch.int32, + device=device, + ) + lengths = torch.tensor([8, 9], dtype=torch.int32, device=device) + out = torch.empty_like(q) + repeat = torch.empty_like(q) + scale = 512**-0.5 + + sm70_sparse_attention_gathered(q, kv, indices, lengths, scale, None, out) + sm70_sparse_attention_gathered(q, kv, indices, lengths, scale, None, repeat) + torch.accelerator.synchronize() + + refs = [] + for row, length in enumerate(lengths.cpu().tolist()): + selected = kv[indices[row, :length].long()].float() + probs = torch.softmax(q[row].float() @ selected.T * scale, dim=-1) + refs.append(probs @ selected) + ref = torch.stack(refs).half() + + torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-3) + assert torch.equal(out, repeat) + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.is_device_capability((7, 0)), + reason="requires an exact SM70 CUDA device", +) +def test_glm53_sm70_fp8_kv_sparse_attention_matches_reference(): + torch.manual_seed(20260827) + device = torch.device("cuda") + num_kv = 48 + block_size = 64 + kv = torch.randn(num_kv, 512, dtype=torch.float16, device=device) * 0.1 + cache = torch.zeros(2, block_size, 520, dtype=torch.uint8, device=device) + slots = torch.arange(num_kv, dtype=torch.int64, device=device) + sm70_glm5_fp8_kv_insert(kv, cache, slots) + + q = torch.randn(2, 16, 512, dtype=torch.float16, device=device) * 0.1 + indices = torch.tensor( + [ + [0, 3, 5, 9, 11, 17, 23, 31] + [-1] * 8, + [2, 4, 8, 12, 18, 24, 32, 40, 45] + [-1] * 7, + ], + dtype=torch.int32, + device=device, + ) + lengths = torch.tensor([8, 9], dtype=torch.int32, device=device) + out = torch.empty_like(q) + repeat = torch.empty_like(q) + scale = 512**-0.5 + sm70_glm5_sparse_attention_paged_fp8(q, cache, indices, lengths, scale, out) + sm70_glm5_sparse_attention_paged_fp8(q, cache, indices, lengths, scale, repeat) + torch.accelerator.synchronize() + + refs = [] + for row, length in enumerate(lengths.cpu().tolist()): + selected = kv[indices[row, :length].long()].float() + probs = torch.softmax(q[row].float() @ selected.T * scale, dim=-1) + refs.append(probs @ selected) + ref = torch.stack(refs).half() + + torch.testing.assert_close(out, ref, rtol=3e-2, atol=4e-3) + assert torch.equal(out, repeat) + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.is_device_capability((7, 0)), + reason="requires an exact SM70 CUDA device", +) +def test_glm53_sm70_fp8_kv_gemm_decode_matches_reference(): + torch.manual_seed(20260828) + device = torch.device("cuda") + num_kv = 48 + index_width = 64 + kv = torch.randn(num_kv, 512, dtype=torch.float16, device=device) * 0.1 + cache = torch.zeros(1, 64, 520, dtype=torch.uint8, device=device) + sm70_glm5_fp8_kv_insert( + kv, + cache, + torch.arange(num_kv, dtype=torch.int64, device=device), + ) + q = torch.randn(1, 16, 512, dtype=torch.float16, device=device) * 0.1 + selected_ids = [0, 3, 5, 9, 11, 17, 23, 31, 40, 45] + indices = torch.tensor( + [selected_ids + [-1] * (index_width - len(selected_ids))], + dtype=torch.int32, + device=device, + ) + lengths = torch.tensor([len(selected_ids)], dtype=torch.int32, device=device) + out = torch.empty_like(q) + repeat = torch.empty_like(q) + gathered_kv = torch.empty(index_width, 512, dtype=torch.float16, device=device) + scores = torch.empty(16, index_width, dtype=torch.float16, device=device) + probs = torch.empty_like(scores) + scale = 512**-0.5 + + sm70_glm5_sparse_attention_paged_fp8_gemm( + q, + cache, + indices, + lengths, + scale, + out, + gathered_kv, + scores, + probs, + ) + sm70_glm5_sparse_attention_paged_fp8_gemm( + q, + cache, + indices, + lengths, + scale, + repeat, + gathered_kv, + scores, + probs, + ) + torch.accelerator.synchronize() + + selected = kv[indices[0, : len(selected_ids)].long()].float() + ref_probs = torch.softmax(q[0].float() @ selected.T * scale, dim=-1) + ref = (ref_probs @ selected).half().unsqueeze(0) + torch.testing.assert_close(out, ref, rtol=3e-2, atol=4e-3) + assert torch.equal(out, repeat) + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.is_device_capability((7, 0)), + reason="requires an exact SM70 CUDA device", +) +def test_glm53_sm70_fp16_fwht_matches_reference(): + torch.manual_seed(11) + q = torch.randn(37, 128, dtype=torch.float16, device="cuda") * 0.1 + + actual = fwht128(q) + repeat = fwht128(q) + reference = _hadamard128_torch(q.float()).half() + torch.accelerator.synchronize() + + torch.testing.assert_close(actual, reference, rtol=1e-2, atol=2e-3) + assert torch.equal(actual, repeat) + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.is_device_capability((7, 0)), + reason="requires an exact SM70 CUDA device", +) +def test_glm53_sm70_fused_q_kv_rmsnorm_matches_reference(): + torch.manual_seed(13) + q = torch.randn(7, 1536, dtype=torch.float16, device="cuda") + kv = torch.randn(7, 512, dtype=torch.float16, device="cuda") + q_weight = torch.randn(1536, dtype=torch.float16, device="cuda") + kv_weight = torch.randn(512, dtype=torch.float16, device="cuda") + eps = 1e-6 + + q_actual, kv_actual = fused_q_kv_rmsnorm(q, kv, q_weight, kv_weight, eps) + q_ref = q.float() * torch.rsqrt(q.float().square().mean(-1, keepdim=True) + eps) + kv_ref = kv.float() * torch.rsqrt(kv.float().square().mean(-1, keepdim=True) + eps) + q_ref = (q_ref * q_weight.float()).half() + kv_ref = (kv_ref * kv_weight.float()).half() + + torch.testing.assert_close(q_actual, q_ref, rtol=2e-3, atol=2e-3) + torch.testing.assert_close(kv_actual, kv_ref, rtol=2e-3, atol=2e-3) diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py index a355501fdd..00475a5925 100644 --- a/tests/models/multimodal/processing/test_gemma4.py +++ b/tests/models/multimodal/processing/test_gemma4.py @@ -10,6 +10,7 @@ from vllm.model_executor.models.gemma4_mm import ( Gemma4ForConditionalGeneration, Gemma4ImagePixelInputs, + _build_broadcast_additive_padding_mask, ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFieldConfig @@ -22,6 +23,51 @@ GEMMA4_MODEL_ID = "google/gemma-4-E2B-it" +def test_broadcast_additive_padding_mask_matches_bool_sdpa() -> None: + torch.manual_seed(0) + padding_positions = torch.tensor( + [[False, False, True, True], [False, False, False, True]] + ) + additive_mask = _build_broadcast_additive_padding_mask( + padding_positions, torch.float32 + ) + + assert additive_mask.shape == (2, 1, 1, 4) + assert additive_mask.dtype == torch.float32 + assert torch.equal(additive_mask.isneginf(), padding_positions[:, None, None, :]) + + query = torch.randn(2, 2, 4, 8) + key = torch.randn(2, 2, 4, 8) + value = torch.randn(2, 2, 4, 8) + bool_mask = (~padding_positions)[:, None, None, :] + + bool_output = torch.nn.functional.scaled_dot_product_attention( + query, key, value, attn_mask=bool_mask + ) + additive_output = torch.nn.functional.scaled_dot_product_attention( + query, key, value, attn_mask=additive_mask + ) + torch.testing.assert_close(additive_output, bool_output, rtol=0, atol=0) + + +def test_broadcast_additive_padding_mask_without_padding_matches_none() -> None: + padding_positions = torch.zeros((1, 4), dtype=torch.bool) + additive_mask = _build_broadcast_additive_padding_mask( + padding_positions, torch.float32 + ) + query = torch.randn(1, 2, 4, 8) + key = torch.randn(1, 2, 4, 8) + value = torch.randn(1, 2, 4, 8) + + unmasked_output = torch.nn.functional.scaled_dot_product_attention( + query, key, value + ) + additive_output = torch.nn.functional.scaled_dot_product_attention( + query, key, value, attn_mask=additive_mask + ) + torch.testing.assert_close(additive_output, unmasked_output, rtol=0, atol=0) + + def test_gemma4_image_schema_accepts_variable_patch_counts(): Gemma4ImagePixelInputs( pixel_values=[ diff --git a/tests/models/qwen4_exp/__init__.py b/tests/models/qwen4_exp/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/models/qwen4_exp/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py new file mode 100644 index 0000000000..c85602a401 --- /dev/null +++ b/tests/models/qwen4_exp/test_config.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm.config import VllmConfig +from vllm.model_executor.models.config import ( + Qwen3_5ForConditionalGenerationConfig, + Qwen4ExpForConditionalGenerationConfig, +) +from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState +from vllm.v1.attention.backends.short_conv_attn import ( + PleShortConvAttentionMetadataBuilder, +) +from vllm.v1.worker.gpu.model_states.mamba_hybrid import ( + MambaHybridAttnMetadata, + MambaHybridModelState, +) + + +def test_sm70_v2_route_accepts_prefix_caching( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + Qwen3_5ForConditionalGenerationConfig, + "verify_and_update_config", + lambda _config: None, + ) + text_config = SimpleNamespace( + hc_count=4, + ple_layer_ids=[2], + indexer_n_heads=4, + rope_parameters={"mrope_section": [11, 11, 10]}, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(rope_parameters={"mrope_interleaved": True}), + hf_text_config=text_config, + multimodal_config=SimpleNamespace(language_model_only=True), + ), + cache_config=SimpleNamespace(enable_prefix_caching=True), + parallel_config=SimpleNamespace(enable_dbo=False, ubatch_size=1), + speculative_config=None, + use_v2_model_runner=True, + ) + + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + assert "mrope_section" not in text_config.rope_parameters + assert "mrope_interleaved" not in vllm_config.model_config.hf_config.rope_parameters + + +def test_initial_sm70_route_rejects_multimodal_tower( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + Qwen3_5ForConditionalGenerationConfig, + "verify_and_update_config", + lambda _config: None, + ) + text_config = SimpleNamespace( + hc_count=4, + ple_layer_ids=[2], + indexer_n_heads=4, + rope_parameters={"mrope_section": [11, 11, 10]}, + ) + multimodal_config = SimpleNamespace(language_model_only=False) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(rope_parameters={}), + hf_text_config=text_config, + multimodal_config=multimodal_config, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False), + parallel_config=SimpleNamespace(enable_dbo=False, ubatch_size=1), + speculative_config=None, + ) + + with pytest.raises(NotImplementedError, match="--language-model-only"): + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + + +@pytest.mark.parametrize( + "architecture", + ["Qwen4ExpForCausalLM", "Qwen4ExpForConditionalGeneration"], +) +def test_qwen4_exp_defaults_to_v2_even_when_quantized_moe( + architecture: str, +) -> None: + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + runner_type="generate", + architectures=[architecture], + is_moe=True, + is_quantized=True, + ) + ) + + assert VllmConfig._is_default_v2_model_runner_model(vllm_config) + + +def test_initial_sm70_v2_route_accepts_native_mtp( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + Qwen3_5ForConditionalGenerationConfig, + "verify_and_update_config", + lambda _config: None, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace( + hc_count=4, + ple_layer_ids=[2], + indexer_n_heads=4, + ), + multimodal_config=None, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False), + parallel_config=SimpleNamespace(enable_dbo=False, ubatch_size=1), + speculative_config=SimpleNamespace(method="mtp"), + ) + + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + + +def test_initial_sm70_v2_route_rejects_unvalidated_speculator( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + Qwen3_5ForConditionalGenerationConfig, + "verify_and_update_config", + lambda _config: None, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace( + hc_count=4, + ple_layer_ids=[2], + indexer_n_heads=4, + ), + multimodal_config=None, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False), + parallel_config=SimpleNamespace(enable_dbo=False, ubatch_size=1), + speculative_config=SimpleNamespace(method="dflash"), + ) + + with pytest.raises(NotImplementedError, match="supports only its native MTP"): + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + + +def test_qwen4_exp_registers_v2_model_state() -> None: + from vllm.models.qwen4_exp.nvidia.model import ( + Qwen4ExpForCausalLM, + Qwen4ExpForConditionalGeneration, + ) + + assert Qwen4ExpForCausalLM.get_model_state_cls() is Qwen4ExpModelState + assert Qwen4ExpForConditionalGeneration.get_model_state_cls() is Qwen4ExpModelState + + +def test_qwen4_exp_v2_model_state_uses_committed_ngram_context() -> None: + model_state = object.__new__(Qwen4ExpModelState) + model_state.uses_ngram_embedding = True + model_state.ngram_context_len = 3 + model_state.ngram_eos_token_id = 99 + model_state.ngram_context = torch.empty((4, 3), dtype=torch.int32) + model_state.ngram_context_offsets = torch.arange(-3, 0, dtype=torch.int64) + model_state.ple_query_start_loc = torch.empty(5, dtype=torch.int32) + + input_batch = SimpleNamespace( + num_reqs=2, + num_reqs_after_padding=3, + idx_mapping=torch.tensor([1, 0]), + query_start_loc=torch.tensor([0, 2, 3, 3], dtype=torch.int32), + ) + req_states = SimpleNamespace( + num_computed_tokens=SimpleNamespace(gpu=torch.tensor([3, 1])), + all_token_ids=SimpleNamespace( + gpu=torch.tensor([[1, 2, 3, 4], [20, 21, 22, 23]], dtype=torch.int32) + ), + ) + + with patch.object(MambaHybridModelState, "prepare_inputs", return_value={}): + model_inputs = model_state.prepare_inputs(input_batch, req_states) + + torch.testing.assert_close( + model_inputs["query_start_loc"], + torch.tensor([0, 2, 3, 3], dtype=torch.int32), + ) + torch.testing.assert_close( + model_inputs["ngram_context"], + torch.tensor([[99, 99, 20], [1, 2, 3], [99, 99, 99]], dtype=torch.int32), + ) + + +def test_qwen4_exp_ple_builder_receives_v2_decode_metadata() -> None: + num_accepted_tokens = torch.tensor([1, 2], dtype=torch.int32) + num_decode_draft_tokens_cpu = torch.tensor([-1, 2], dtype=torch.int32) + metadata = MambaHybridAttnMetadata( + is_prefilling=torch.tensor([False, False]), + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) + builder = PleShortConvAttentionMetadataBuilder.__new__( + PleShortConvAttentionMetadataBuilder + ) + + kwargs = metadata.get_extra_attn_kwargs(builder, num_reqs=2) + + torch.testing.assert_close(kwargs["num_accepted_tokens"], num_accepted_tokens) + torch.testing.assert_close( + kwargs["num_decode_draft_tokens_cpu"], num_decode_draft_tokens_cpu + ) + + +def test_qwen4_exp_qsa_owns_qkv_projection_for_private_qwen3_api() -> None: + from vllm.models.qwen4_exp.nvidia.qsa import Qwen4ExpQSAAttention + + class AddRotaryOffset(torch.nn.Module): + def forward(self, positions, query, key): + del positions + return query + 1, key + 2 + + layer = object.__new__(Qwen4ExpQSAAttention) + torch.nn.Module.__init__(layer) + layer.q_size = 4 + layer.kv_size = 2 + layer.num_heads = 2 + layer.num_kv_heads = 1 + layer.head_dim = 2 + layer.q_norm = torch.nn.Identity() + layer.k_norm = torch.nn.Identity() + layer.rotary_emb = AddRotaryOffset() + qkv = torch.arange(12, dtype=torch.float32).reshape(1, 12) + + query, key, value, gate = layer._project_qkv_gate( + qkv, torch.tensor([0], dtype=torch.int64) + ) + + torch.testing.assert_close(query, torch.tensor([[1, 2, 5, 6.0]])) + torch.testing.assert_close(gate, torch.tensor([[2, 3, 6, 7.0]])) + torch.testing.assert_close(key, torch.tensor([[10, 11.0]])) + torch.testing.assert_close(value, torch.tensor([[10, 11.0]])) + + +def test_qwen4_exp_qsa_uses_private_mrope_signature( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.models.qwen4_exp.nvidia import indexer_qsa + + calls = [] + + def private_triton_mrope( + query, + key, + cos, + sin, + mrope_section, + head_size, + rotary_dim, + mrope_interleaved, + ): + calls.append( + ( + cos.shape, + sin.shape, + mrope_section, + head_size, + rotary_dim, + mrope_interleaved, + ) + ) + return query, key + + monkeypatch.setattr(indexer_qsa, "triton_mrope", private_triton_mrope) + cache = torch.zeros(16, 4) + rotary_emb = SimpleNamespace( + rotary_dim=2, + mrope_section=[1, 0, 0], + mrope_interleaved=True, + _match_cos_sin_cache_dtype=lambda _tensor: cache, + ) + tensor = torch.arange(8, dtype=torch.float32).reshape(2, 1, 4) + positions = torch.tensor([[0, 1], [2, 3], [4, 5]]) + + output = indexer_qsa.apply_qsa_rope(rotary_emb, positions, tensor) + + torch.testing.assert_close(output, tensor) + assert calls == [ + (torch.Size([3, 2, 2]), torch.Size([3, 2, 2]), [1, 0, 0], 4, 2, True) + ] diff --git a/tests/models/qwen4_exp/test_model_forward.py b/tests/models/qwen4_exp/test_model_forward.py new file mode 100644 index 0000000000..9a34e3cf31 --- /dev/null +++ b/tests/models/qwen4_exp/test_model_forward.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer + + +class _AttentionHyperConnection: + def mix(self, hidden_states: torch.Tensor): + return hidden_states, hidden_states, torch.zeros_like(hidden_states) + + +class _MlpHyperConnection: + def combine_and_mix( + self, + hidden_states: torch.Tensor, + block_output: torch.Tensor, + injection: torch.Tensor, + ): + return hidden_states, block_output, injection + + +class _OutputOnlyGdn: + def __init__(self) -> None: + self.output: torch.Tensor | None = None + + def __call__( + self, + *, + hidden_states: torch.Tensor, + output: torch.Tensor | None, + ) -> torch.Tensor: + assert output is not None + self.output = output + output.copy_(hidden_states + 1) + return hidden_states + 100 + + +def test_linear_attention_forwards_preallocated_output_buffer() -> None: + layer = object.__new__(Qwen4ExpDecoderLayer) + object.__setattr__(layer, "ple", None) + object.__setattr__(layer, "layer_type", "linear_attention") + object.__setattr__(layer, "attn_hyper_connection", _AttentionHyperConnection()) + object.__setattr__(layer, "mlp_hyper_connection", _MlpHyperConnection()) + gdn = _OutputOnlyGdn() + object.__setattr__(layer, "linear_attn", gdn) + object.__setattr__(layer, "mlp", lambda hidden_states: hidden_states) + hidden_states = torch.arange(6, dtype=torch.float32).view(2, 3) + + _, mlp_out, _ = Qwen4ExpDecoderLayer.forward( + layer, + hidden_states, + None, + None, + torch.arange(2), + input_ids=None, + query_start_loc=None, + ngram_context=None, + ) + + assert gdn.output is not None + torch.testing.assert_close(mlp_out, hidden_states + 1) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py new file mode 100644 index 0000000000..6777828c76 --- /dev/null +++ b/tests/models/qwen4_exp/test_ple.py @@ -0,0 +1,495 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from torch.nn import functional as F + +import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module +import vllm.model_executor.parameter as parameter_module +import vllm.models.qwen4_exp.nvidia.ple_layer as ple_module +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.model_loader.utils import device_loading_context +from vllm.models.qwen4_exp.common.ple import ( + PLEShardOverlap, + compute_ple_shard_overlap, + copy_ple_embedding_shard_, +) +from vllm.models.qwen4_exp.nvidia.ple_layer import ( + Qwen4ExpNGramEmbedding, + Qwen4ExpPinnedHostEmbedding, + Qwen4ExpPLEFp8EmbeddingMethod, + Qwen4ExpPLELayer, + _get_ple_embedding_quant_method, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA pinned memory") +def test_pinned_host_ple_allocates_tp_shard_without_device_table( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(ple_module, "is_pin_memory_available", lambda: True) + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 2) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 4 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 2) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 4 + ) + + layer = Qwen4ExpPinnedHostEmbedding( + num_embeddings=32, + embedding_dim=8, + params_dtype=torch.float16, + padding_size=8, + prefix="model.layers.2.ple.ngram_embedding", + quant_method=Qwen4ExpPLEFp8EmbeddingMethod(), + ) + + assert layer.tp_size == 4 + assert layer.weight.shape == (8, 8) + assert layer.weight.dtype == torch.float8_e4m3fn + assert layer.weight.device.type == "cpu" + assert layer.weight.is_pinned() + assert layer.weight._vllm_keep_on_cpu + assert not layer.weight.is_meta + assert not layer.weight_scale.is_meta + assert layer.weight_scale.dtype == torch.float16 + assert layer._accelerator_weight_views == {} + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), + reason="requires an exact SM70 CUDA device", +) +def test_pinned_host_ple_fp8_rows_are_gatherable_on_sm70( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(ple_module, "is_pin_memory_available", lambda: True) + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + embedding_module, "tensor_model_parallel_all_reduce", lambda tensor: tensor + ) + layer = Qwen4ExpPinnedHostEmbedding( + num_embeddings=8, + embedding_dim=8, + params_dtype=torch.float16, + padding_size=8, + prefix="model.layers.2.ple.ngram_embedding", + quant_method=Qwen4ExpPLEFp8EmbeddingMethod(), + ) + raw = torch.tensor( + [0x00, 0x01, 0x08, 0x38, 0x7E, 0x80, 0xB8, 0xFE], + dtype=torch.uint8, + ).repeat(8, 1) + layer.weight.data.copy_(raw.view(torch.float8_e4m3fn)) + layer.weight_scale = nn.Parameter( + torch.tensor([0.25], dtype=torch.float16, device="cuda"), + requires_grad=False, + ) + layer.prepare_accelerator_weight() + + output = layer(torch.tensor([0, 7], dtype=torch.int64, device="cuda")) + torch.accelerator.synchronize() + + assert output.dtype == torch.float16 + expected = raw.view(torch.float8_e4m3fn).float() * 0.25 + torch.testing.assert_close(output.float().cpu(), expected[[0, 7]]) + + +def test_post_load_context_keeps_marked_parameter_on_cpu() -> None: + module = nn.Module() + host_weight = nn.Parameter(torch.ones(2)) + host_weight._vllm_keep_on_cpu = True + module.register_parameter("weight", host_weight) + + with device_loading_context(module, torch.device("meta")): + assert module.weight.device.type == "cpu" + + assert module.weight.device.type == "cpu" + + +def test_ngram_embedding_accepts_checkpoint_seed_none( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 4 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 4 + ) + config = SimpleNamespace( + ngram_size=3, + heads_per_ngram=8, + eos_token_id=2, + vocab_size=64, + split_ngram_parts=2, + seed=None, + ngram_vocab_size_base=101, + make_ngram_vocab_size_divisible_by=128, + ple_embedding_dtype="float8_e4m3fn", + ple_offload_embedding=False, + ) + + with torch.device("meta"): + layer = Qwen4ExpNGramEmbedding( + config, + embedding_dim=256, + ple_dense_layer_id=0, + max_total_tokens=8, + max_num_reqs=2, + prefix="model.layers.2.ple.ple_embedding", + params_dtype=torch.float16, + ) + + assert layer.ngram_heads == 16 + assert layer.head_dim == 16 + assert layer.ngram_embedding.weight.dtype == torch.float8_e4m3fn + assert layer.ngram_embedding.weight.is_meta + + +def _make_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding: + module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding) + nn.Module.__init__(module) + module.split_ngram_parts = 2 + module.register_buffer("layer_multipliers", torch.zeros(1, dtype=torch.long)) + module.register_buffer("ngram_heads_offsets", torch.zeros(1, dtype=torch.long)) + module.register_buffer("ngram_heads_vocab_sizes", torch.zeros(1, dtype=torch.long)) + module.ngram_embedding = SimpleNamespace( + org_vocab_size=8, + embedding_dim=2, + weight=nn.Parameter(torch.full((4, 2), -1.0)), + shard_indices=SimpleNamespace( + org_vocab_start_index=2, + org_vocab_end_index=6, + ), + ) + return module + + +def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding: + module = _make_ngram_embedding_for_load_test() + embedding = nn.Module() + embedding.org_vocab_size = 8 + embedding.embedding_dim = 2 + embedding.shard_indices = SimpleNamespace( + org_vocab_start_index=2, + org_vocab_end_index=6, + ) + embedding.register_parameter( + "weight", + nn.Parameter( + torch.full((4, 2), -1.0).to(torch.float8_e4m3fn), + requires_grad=False, + ), + ) + embedding.register_parameter( + "weight_scale", + nn.Parameter(torch.zeros(1, dtype=torch.bfloat16), requires_grad=False), + ) + module.ngram_embedding = embedding + return module + + +def test_ple_shard_overlap_and_copy() -> None: + overlap = compute_ple_shard_overlap( + checkpoint_start=2, checkpoint_rows=5, tp_start=4, tp_end=8 + ) + assert overlap == PLEShardOverlap(source_start=2, destination_start=0, row_count=3) + + destination = torch.full((4, 2), -1.0) + loaded = torch.arange(10, dtype=torch.float64).reshape(5, 2) + copied = copy_ple_embedding_shard_( + destination, + loaded, + checkpoint_start=2, + tp_start=4, + tp_end=8, + ) + + assert copied == 3 + torch.testing.assert_close(destination[:3], loaded[2:5].float()) + torch.testing.assert_close(destination[3], torch.tensor([-1.0, -1.0])) + + +def test_ple_shard_copy_is_a_noop_without_overlap() -> None: + destination = torch.ones(4, 2) + copied = copy_ple_embedding_shard_( + destination, + torch.zeros(2, 2), + checkpoint_start=10, + tp_start=4, + tp_end=8, + ) + + assert copied == 0 + assert torch.equal(destination, torch.ones_like(destination)) + + +def test_ngram_embedding_loads_shards_and_ignores_legacy_token_lookup() -> None: + module = _make_ngram_embedding_for_load_test() + shard_0 = torch.arange(8, dtype=torch.float32).reshape(4, 2) + shard_1 = torch.arange(8, 16, dtype=torch.float32).reshape(4, 2) + + loaded = module.load_weights( + [ + ("ngram_embedding.shard_0.weight", shard_0), + ("ngram_embedding.shard_1.weight", shard_1), + ("token_lookup", torch.tensor([2, 1, 0])), + ] + ) + + assert loaded == {"ngram_embedding.weight"} + torch.testing.assert_close( + module.ngram_embedding.weight, + torch.cat((shard_0[2:4], shard_1[0:2])), + ) + + +def test_ngram_embedding_rejects_mismatched_checkpoint_shard() -> None: + module = _make_ngram_embedding_for_load_test() + + with pytest.raises( + ValueError, + match=r"Shape mismatch for PLE embedding shard 0", + ): + module.load_weights([("ngram_embedding.shard_0.weight", torch.zeros(3, 2))]) + + +def test_ngram_embedding_loads_fp8_shards_and_global_scale() -> None: + module = _make_fp8_ngram_embedding_for_load_test() + shard_0 = torch.arange(8, dtype=torch.float32).reshape(4, 2).to(torch.float8_e4m3fn) + shard_1 = ( + torch.arange(8, 16, dtype=torch.float32).reshape(4, 2).to(torch.float8_e4m3fn) + ) + weight_scale = torch.tensor([0.25], dtype=torch.bfloat16) + + loaded = module.load_weights( + [ + ("ngram_embedding.shard_0.weight", shard_0), + ("ngram_embedding.shard_1.weight", shard_1), + ("ngram_embedding.weight_scale", weight_scale), + ] + ) + + assert loaded == {"ngram_embedding.weight", "ngram_embedding.weight_scale"} + assert module.ngram_embedding.weight.dtype == torch.float8_e4m3fn + assert torch.equal( + module.ngram_embedding.weight.float(), + torch.cat((shard_0[2:4], shard_1[0:2])).float(), + ) + assert torch.equal(module.ngram_embedding.weight_scale, weight_scale) + + +def _make_fp8_embedding_layer( + monkeypatch: pytest.MonkeyPatch, + params_dtype: torch.dtype = torch.bfloat16, +) -> embedding_module.VocabParallelEmbedding: + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + embedding_module, + "tensor_model_parallel_all_reduce", + lambda tensor: tensor, + ) + method = Qwen4ExpPLEFp8EmbeddingMethod() + layer = embedding_module.VocabParallelEmbedding( + 3, + 2, + params_dtype=params_dtype, + padding_size=1, + quant_method=method, + ) + weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]]) + layer.weight.data.copy_(weight.to(torch.float8_e4m3fn)) + layer.weight_scale.data.copy_(torch.tensor([0.25], dtype=params_dtype)) + return layer + + +def test_ple_fp8_embedding_scale_matches_model_dtype(monkeypatch) -> None: + layer = _make_fp8_embedding_layer(monkeypatch, params_dtype=torch.float16) + + assert layer.weight_scale.dtype == torch.float16 + + +def test_ple_fp8_embedding_dequantizes_in_ple_layer(monkeypatch) -> None: + layer = _make_fp8_embedding_layer(monkeypatch) + quantized_output = layer(torch.tensor([2, 0])) + ple_layer = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer) + nn.Module.__init__(ple_layer) + ple_layer.ple_embedding = nn.Module() + ple_layer.ple_embedding.ngram_embedding = layer + + output = ple_layer._dequantize_embeddings( + quantized_output, + torch.bfloat16, + ) + + assert layer.weight.dtype == torch.float8_e4m3fn + assert layer.weight_scale.dtype == torch.bfloat16 + assert quantized_output.dtype == torch.float8_e4m3fn + assert output.dtype == torch.bfloat16 + weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]]) + torch.testing.assert_close(output, (weight[[2, 0]] * 0.25).bfloat16()) + + +def test_ple_fp8_embedding_uses_int8_for_tp_reduce(monkeypatch) -> None: + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 2 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 2 + ) + monkeypatch.setattr( + embedding_module, + "get_masked_input_and_mask", + lambda *args: ( + torch.tensor([0, 0]), + torch.tensor([False, True]), + ), + ) + reduced_dtypes = [] + + def all_reduce(tensor: torch.Tensor) -> torch.Tensor: + reduced_dtypes.append(tensor.dtype) + return tensor.clone() + + monkeypatch.setattr( + embedding_module, + "tensor_model_parallel_all_reduce", + all_reduce, + ) + layer = embedding_module.VocabParallelEmbedding( + 4, + 2, + params_dtype=torch.bfloat16, + padding_size=1, + quant_method=Qwen4ExpPLEFp8EmbeddingMethod(), + ) + layer.weight.data.copy_( + torch.tensor([[1.0, 2.0], [4.0, 8.0]]).to(torch.float8_e4m3fn) + ) + + output = layer(torch.tensor([0, 2])) + + assert reduced_dtypes == [torch.int8] + assert output.dtype == torch.float8_e4m3fn + torch.testing.assert_close(output[0].float(), layer.weight[0].float()) + assert torch.count_nonzero(output[1].float()) == 0 + + +def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None: + prefix = "model.layers.1.ple.ple_embedding.ngram_embedding" + quant_config = Fp8Config( + is_checkpoint_fp8_serialized=True, + ignored_layers=[], + weight_block_size=[128, 128], + ) + assert isinstance( + _get_ple_embedding_quant_method(quant_config, prefix), + Qwen4ExpPLEFp8EmbeddingMethod, + ) + + quant_config.ignored_layers = [f"{prefix}.shard_0"] + assert _get_ple_embedding_quant_method(quant_config, prefix) is None + + +def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None: + module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer) + nn.Module.__init__(module) + module.conv_state_len = 6 + module.short_conv_dilation = 2 + + conv_weights = torch.tensor([[0.25, -0.5, 0.75, 1.0]]) + conv_state = torch.zeros(2, 1, 9) + conv_state[1] = torch.arange(1, 10, dtype=torch.float32).reshape(1, 9) + first_inputs = torch.tensor([[10.0], [20.0], [30.0], [40.0]]) + initial_state = conv_state[1:].clone() + first_history = torch.cat( + (initial_state[..., : module.conv_state_len], first_inputs.T.unsqueeze(0)), + dim=-1, + ) + + graph_padded_inputs = F.pad(first_inputs, (0, 0, 0, 4)) + first_output = module._short_conv_dilated_spec_batched( + graph_padded_inputs, + conv_state, + conv_weights, + torch.tensor([1, 0]), + torch.tensor([0, 4, 4]), + torch.tensor([1, 0]), + spec_query_len=4, + ) + + expected_first_output = F.silu( + F.conv1d( + first_history, + conv_weights.unsqueeze(1), + groups=1, + dilation=module.short_conv_dilation, + ) + ).transpose(1, 2)[0] + expected_first_state = first_history[..., 1:10] + torch.testing.assert_close(first_output[:4], expected_first_output) + assert torch.count_nonzero(first_output[4:]) == 0 + assert torch.count_nonzero(conv_state[0]) == 0 + torch.testing.assert_close(conv_state[1:], expected_first_state) + + second_inputs = torch.tensor([[50.0], [60.0]]) + rollback_state = expected_first_state[..., 1:7] + padded_second_inputs = F.pad(second_inputs.T.unsqueeze(0), (0, 2)) + second_history = torch.cat((rollback_state, padded_second_inputs), dim=-1) + expected_second_state = expected_first_state.clone() + expected_second_state[..., :7] = second_history[..., 1:8] + + second_output = module._short_conv_dilated_spec_batched( + second_inputs, + conv_state, + conv_weights, + torch.tensor([1]), + torch.tensor([0, 2]), + torch.tensor([2]), + spec_query_len=4, + ) + + expected_second_output = F.silu( + F.conv1d( + second_history, + conv_weights.unsqueeze(1), + groups=1, + dilation=module.short_conv_dilation, + ) + ).transpose(1, 2)[0, :2] + torch.testing.assert_close(second_output, expected_second_output) + torch.testing.assert_close(conv_state[1:], expected_second_state) + + +def test_ple_state_shape_reserves_speculative_tokens() -> None: + module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer) + nn.Module.__init__(module) + module.hc_hidden_size = 32 + module.conv_state_len = 9 + module.num_spec_tokens = 3 + + assert module.get_state_shape()[0] in ((32, 12), (12, 32)) diff --git a/tests/models/qwen4_exp/test_qsa_cache.py b/tests/models/qwen4_exp/test_qsa_cache.py new file mode 100644 index 0000000000..d3cfea8d5f --- /dev/null +++ b/tests/models/qwen4_exp/test_qsa_cache.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from typing import Any + +import torch + +from vllm.models.qwen4_exp.common.qsa_cache import QSAKeyStateCache +from vllm.models.qwen4_exp.nvidia.qsa import Qwen4ExpQSAFlashAttentionImpl +from vllm.v1.worker.utils import bind_kv_cache + + +def test_bind_qsa_key_cache_builds_key_and_mrope_views() -> None: + prefix = "model.layers.0.self_attn.raw_key_cache" + static_forward_context: dict[str, Any] = {} + layer = QSAKeyStateCache( + head_size=128, + dtype=torch.float16, + cache_rope_positions=True, + prefix=prefix, + cache_config=SimpleNamespace(block_size=16), + compress_ratio=4, + vllm_config=SimpleNamespace( + compilation_config=SimpleNamespace( + static_forward_context=static_forward_context + ) + ), + ) + cache = torch.empty(2, 8, 1, layer.head_size, dtype=torch.float16) + runner_kv_caches: list[torch.Tensor] = [] + + bind_kv_cache({prefix: cache}, static_forward_context, runner_kv_caches) + + assert layer.kv_cache is cache + assert layer.key_cache.shape == (2, 8, 1, 128) + assert layer.key_cache.untyped_storage().data_ptr() == cache.data_ptr() + assert layer.rope_position_cache.shape == (2, 8, 1, 3) + assert layer.rope_position_cache.dtype == torch.int64 + assert runner_kv_caches == [cache] + + +def test_qsa_forward_splits_local_flash_cache_layout(monkeypatch) -> None: + from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops + + num_blocks, block_size, head_size = 2, 16, 8 + kv_cache = torch.arange( + num_blocks * 2 * block_size * head_size, + dtype=torch.float16, + ).view(num_blocks, 2, block_size, 1, head_size) + query = torch.zeros(1, 2, head_size, dtype=torch.float16) + output = torch.empty_like(query) + logical_indices = torch.zeros(1, 4, dtype=torch.int32) + block_table = torch.zeros(1, 1, dtype=torch.int32) + token_to_req = torch.zeros(1, dtype=torch.int32) + captured = {} + + def fake_sparse_attention( + query_arg, + key_cache_arg, + value_cache_arg, + logical_indices_arg, + block_table_arg, + token_to_req_arg, + output_arg, + ): + captured["key_cache"] = key_cache_arg + captured["value_cache"] = value_cache_arg + assert torch.equal(query_arg, query) + assert torch.equal(logical_indices_arg, logical_indices) + assert torch.equal(block_table_arg, block_table) + assert torch.equal(token_to_req_arg, token_to_req) + output_arg.fill_(1) + return output_arg + + monkeypatch.setattr(qsa_ops, "qsa_sparse_paged_attention", fake_sparse_attention) + impl = object.__new__(Qwen4ExpQSAFlashAttentionImpl) + impl.head_size = head_size + impl.alibi_slopes = None + impl.sinks = None + impl.sliding_window = (-1, -1) + + result = impl.forward_qsa( + SimpleNamespace(topk_indices_buffer=logical_indices), + query, + query[:, :1], + query[:, :1], + kv_cache, + SimpleNamespace(num_actual_tokens=1, block_table=block_table), + output, + token_to_req, + ) + + expected_key, expected_value = kv_cache.unbind(1) + assert torch.equal(captured["key_cache"], expected_key) + assert torch.equal(captured["value_cache"], expected_value) + assert result is output + assert torch.equal(output, torch.ones_like(output)) diff --git a/tests/models/qwen4_exp/test_qsa_ops.py b/tests/models/qwen4_exp/test_qsa_ops.py new file mode 100644 index 0000000000..ebc3f3273d --- /dev/null +++ b/tests/models/qwen4_exp/test_qsa_ops.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops +from vllm.models.qwen4_exp.nvidia.ops.qsa import ( + _qsa_indexer_cublas_shape_supported, + _qsa_sparse_launch_profile, +) + + +def test_sm70_qsa_prefill_uses_narrow_tiles_and_four_warps(): + assert _qsa_sparse_launch_profile(511, 8, True) == (64, 4, 4) + assert _qsa_sparse_launch_profile(512, 8, True) == (32, 4, 4) + assert _qsa_sparse_launch_profile(8192, 8, True) == (32, 1, 4) + + +def test_non_sm70_qsa_prefill_keeps_gb300_profile(): + assert _qsa_sparse_launch_profile(512, 8, False) == (64, 4, 2) + assert _qsa_sparse_launch_profile(8192, 8, False) == (64, 1, 2) + + +def test_qsa_indexer_cublas_accepts_only_exact_single_request_shape(): + query = torch.empty(8, 4, 128, dtype=torch.float16) + cache = torch.empty(2, 400, 1, 128, dtype=torch.float16) + page_table = torch.empty(1, 2, dtype=torch.int32) + + assert _qsa_indexer_cublas_shape_supported(query, cache, page_table) + assert not _qsa_indexer_cublas_shape_supported( + query.to(torch.bfloat16), cache, page_table + ) + assert not _qsa_indexer_cublas_shape_supported( + query, cache, page_table.expand(2, -1) + ) + assert not _qsa_indexer_cublas_shape_supported(query[:, :3], cache, page_table) + + +def test_qsa_indexer_cublas_does_not_capture_decode_rows(monkeypatch): + cache = torch.empty(2, 400, 1, 128, dtype=torch.float16) + page_table = torch.empty(1, 2, dtype=torch.int32) + monkeypatch.setattr(qsa_ops, "_SM70_INDEXER_CUBLAS", True) + monkeypatch.setattr(qsa_ops, "_SM70_INDEXER_CUBLAS_MIN_ROWS", 256) + monkeypatch.setattr( + qsa_ops.current_platform, + "is_device_capability", + lambda capability: capability == 70, + ) + + assert qsa_ops._use_sm70_qsa_indexer_cublas( + torch.empty(256, 4, 128, dtype=torch.float16), cache, page_table + ) + assert not qsa_ops._use_sm70_qsa_indexer_cublas( + torch.empty(255, 4, 128, dtype=torch.float16), cache, page_table + ) + + +def test_qsa_indexer_cublas_requires_enough_score_work(monkeypatch): + monkeypatch.setattr( + qsa_ops, + "_SM70_INDEXER_CUBLAS_MIN_SCORE_ELEMENTS", + 1024**2, + ) + + assert not qsa_ops._qsa_indexer_cublas_work_supported(1024, 512) + assert qsa_ops._qsa_indexer_cublas_work_supported(2048, 512) diff --git a/tests/models/qwen4_exp/test_weight_loading.py b/tests/models/qwen4_exp/test_weight_loading.py new file mode 100644 index 0000000000..ce95ac3441 --- /dev/null +++ b/tests/models/qwen4_exp/test_weight_loading.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +from torch import nn + +from vllm.model_executor.layers.fused_moe.layer import FusedMoE +from vllm.model_executor.models.qwen3_next import Qwen3NextSparseMoeBlock +from vllm.models.qwen4_exp.nvidia.model import ( + Qwen4ExpForConditionalGeneration, + Qwen4ExpModel, + Qwen4ExpSparseMoeBlock, + _remap_qsa_cache_scale_name, +) +from vllm.models.qwen4_exp.nvidia.mtp import _validate_mtp_expert_weights_loaded + + +@pytest.mark.parametrize( + ("checkpoint_name", "model_name", "shard_id"), + [ + ( + "layers.0.self_attn.q_proj.weight", + "layers.0.self_attn.qkv_proj.weight", + "q", + ), + ( + "layers.0.self_attn.k_proj.weight", + "layers.0.self_attn.qkv_proj.weight", + "k", + ), + ( + "layers.1.linear_attn.in_proj_qkv.weight", + "layers.1.linear_attn.in_proj_qkvz.weight", + (0, 1, 2), + ), + ( + "layers.1.linear_attn.in_proj_z.weight", + "layers.1.linear_attn.in_proj_qkvz.weight", + 3, + ), + ( + "layers.1.linear_attn.in_proj_b.weight", + "layers.1.linear_attn.in_proj_ba.weight", + 0, + ), + ( + "layers.1.mlp.gate_proj.weight", + "layers.1.mlp.gate_up_proj.weight", + 0, + ), + ( + "layers.1.mlp.experts.0.gate_proj.weight", + "layers.1.mlp.experts.0.gate_proj.weight", + None, + ), + ( + "layers.0.self_attn.indexer.index_qk_proj.weight", + "layers.0.self_attn.indexer.index_qk_proj.weight", + None, + ), + ( + "layers.0.attn_hyper_connection.input_mix_weight_down.weight", + "layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight", + 0, + ), + ( + "layers.0.attn_hyper_connection.block_inject_weight.weight", + "layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight", + 1, + ), + ( + "hyper_connection_mixer.input_mix_weight_down.weight", + "hyper_connection_mixer.input_mix_weight_down.weight", + None, + ), + ( + "layers.1.ple.ple_embedding.layer_multipliers", + "layers.1.ple.ple_embedding.layer_multipliers", + None, + ), + ], +) +def test_text_checkpoint_mapper_preserves_qwen4_exp_specific_weights( + checkpoint_name: str, + model_name: str, + shard_id: str | int | tuple[int, ...] | None, +) -> None: + assert Qwen4ExpModel.hf_to_vllm_mapper._map_name_with_shard(checkpoint_name) == ( + model_name, + shard_id, + ) + + +def test_outer_checkpoint_mapper_selects_language_model_only_paths() -> None: + mapper = Qwen4ExpForConditionalGeneration.hf_to_vllm_mapper + + assert ( + mapper._map_name("model.language_model.layers.0.ple.key_proj.weight") + == "language_model.model.layers.0.ple.key_proj.weight" + ) + assert mapper._map_name("lm_head.weight") == "language_model.lm_head.weight" + assert mapper._map_name("model.visual.blocks.0.attn.qkv.weight") == ( + "visual.blocks.0.attn.qkv.weight" + ) + + +def test_sparse_moe_attaches_private_recursive_loader_mapping(monkeypatch) -> None: + class FakeExperts(nn.Module): + def __init__(self) -> None: + super().__init__() + self.expert_mapping = None + + def fake_qwen3_next_init(self, vllm_config, prefix="") -> None: + del vllm_config, prefix + nn.Module.__init__(self) + self.n_routed_experts = 2 + self.n_redundant_experts = 0 + self.experts = FakeExperts() + + monkeypatch.setattr( + Qwen3NextSparseMoeBlock, + "__init__", + fake_qwen3_next_init, + ) + vllm_config = type( + "VllmConfigStub", + (), + { + "parallel_config": type( + "ParallelConfigStub", (), {"use_sequence_parallel_moe": False} + )(), + "model_config": type( + "ModelConfigStub", + (), + { + "hf_text_config": type( + "TextConfigStub", + (), + {"shared_expert_intermediate_size": 640}, + )() + }, + )(), + }, + )() + + block = Qwen4ExpSparseMoeBlock(vllm_config, prefix="model.layers.0.mlp") + + assert block.experts.expert_mapping == [ + ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), + ("experts.w13_weight", "experts.gate_up_proj", 1, "w3"), + ("experts.w2_weight", "experts.down_proj", 0, "w2"), + ("experts.w13_", "experts.0.gate_proj.", 0, "w1"), + ("experts.w2_", "experts.0.down_proj.", 0, "w2"), + ("experts.w13_", "experts.0.up_proj.", 0, "w3"), + ("experts.w13_", "experts.1.gate_proj.", 1, "w1"), + ("experts.w2_", "experts.1.down_proj.", 1, "w2"), + ("experts.w13_", "experts.1.up_proj.", 1, "w3"), + ] + + +def test_fused_mtp_expert_checkpoint_loads_every_expert() -> None: + class FakeFusedExperts(nn.Module): + load_weights = FusedMoE.load_weights + + def __init__(self) -> None: + super().__init__() + self.layer_name = "model.layers.0.mlp.experts" + self.w13_weight = nn.Parameter(torch.empty(1)) + self.w2_weight = nn.Parameter(torch.empty(1)) + self.calls: list[tuple[str, str, int, torch.Tensor]] = [] + self.expert_mapping = FusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=2, + include_fused=True, + ) + + def weight_loader( + self, + *, + param: nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + return_success: bool, + ) -> bool: + assert return_success + param_name = "w13" if param is self.w13_weight else "w2" + self.calls.append((param_name, shard_id, expert_id, loaded_weight.clone())) + return True + + experts = FakeFusedExperts() + gate_up = torch.arange(2 * 6 * 2).reshape(2, 6, 2) + down = torch.arange(2 * 2 * 3).reshape(2, 2, 3) + + loaded = list( + experts.load_weights([("gate_up_proj", gate_up), ("down_proj", down)]) + ) + + assert loaded == [ + "w13_weight", + "w13_weight", + "w13_weight", + "w13_weight", + "w2_weight", + "w2_weight", + ] + assert [(name, shard, expert) for name, shard, expert, _ in experts.calls] == [ + ("w13", "w1", 0), + ("w13", "w1", 1), + ("w13", "w3", 0), + ("w13", "w3", 1), + ("w2", "w2", 0), + ("w2", "w2", 1), + ] + torch.testing.assert_close(experts.calls[0][3], gate_up[0, :3]) + torch.testing.assert_close(experts.calls[1][3], gate_up[1, :3]) + torch.testing.assert_close(experts.calls[2][3], gate_up[0, 3:]) + torch.testing.assert_close(experts.calls[3][3], gate_up[1, 3:]) + torch.testing.assert_close(experts.calls[4][3], down[0]) + torch.testing.assert_close(experts.calls[5][3], down[1]) + + +def test_mtp_expert_loading_fails_closed() -> None: + model = nn.Module() + model.model = nn.Module() + model.model.layers = nn.ModuleList([nn.Module()]) + model.model.layers[0].mlp = nn.Module() + model.model.layers[0].mlp.experts = nn.Module() + experts = model.model.layers[0].mlp.experts + experts.w13_weight = nn.Parameter(torch.empty(1)) + experts.w2_weight = nn.Parameter(torch.empty(1)) + + w13_name = "model.layers.0.mlp.experts.w13_weight" + w2_name = "model.layers.0.mlp.experts.w2_weight" + _validate_mtp_expert_weights_loaded(model, {w13_name, w2_name}) + + with pytest.raises(ValueError, match="w2_weight"): + _validate_mtp_expert_weights_loaded(model, {w13_name}) + + +@pytest.mark.parametrize( + ("checkpoint_name", "model_name"), + [ + ( + "layers.0.self_attn.k_proj.k_scale", + "layers.0.self_attn._k_scale", + ), + ( + "layers.0.self_attn.v_proj.output_scale", + "layers.0.self_attn._v_scale", + ), + ( + "language_model.model.layers.0.self_attn.attn.k_scale", + "language_model.model.layers.0.self_attn._k_scale", + ), + ( + "layers.0.self_attn.indexer.index_qk_proj.weight_scale", + "layers.0.self_attn.indexer.index_qk_proj.weight_scale", + ), + ( + "layers.1.self_attn.k_proj.k_scale", + "layers.1.self_attn.k_proj.k_scale", + ), + ], +) +def test_only_qsa_main_cache_scales_move_to_the_merged_owner( + checkpoint_name: str, + model_name: str, +) -> None: + assert _remap_qsa_cache_scale_name(checkpoint_name, frozenset({0})) == model_name diff --git a/tests/quantization/test_sm70_fp8_prefill_exact_dense.py b/tests/quantization/test_sm70_fp8_prefill_exact_dense.py index e0a1fb6eaa..166f871954 100644 --- a/tests/quantization/test_sm70_fp8_prefill_exact_dense.py +++ b/tests/quantization/test_sm70_fp8_prefill_exact_dense.py @@ -27,10 +27,12 @@ _get_sm70_fp8_prefill_exact_dense_workspace, _is_sm70_fp8_exact_8k_prefill_layer, _is_sm70_fp8_prefill_exact_dense_layer, + _is_sm70_fp8_prescaled_m1_decode_layer, _is_sm70_fp8_qpn8_layer, _is_sm70_fp8_qpn8_runtime_contract, _sm70_fp8_prefill_dense_workspaces, _sm70_fp8_prefill_visible_dense_mm, + _try_sm70_fp8_prescaled_decode_scales, ) @@ -676,6 +678,101 @@ def fake_prescaled(out, input, qweight, scales, *args): assert calls[2][2] is normal_scales +def test_fp8_prescaled_m1_decode_admission_is_exact(): + layer = SimpleNamespace( + prefix="model.layers.2.attn.fused_wqa_wkv", + tp_size=1, + weight_block_size=[128, 128], + input_size_per_partition=4096, + output_size_per_partition=1536, + weight=torch.empty((1536, 4096), device="meta"), + ) + assert _is_sm70_fp8_prescaled_m1_decode_layer(layer) + + layer.tp_size = 4 + assert not _is_sm70_fp8_prescaled_m1_decode_layer(layer) + layer.tp_size = 1 + layer.prefix = "model.layers.2.attn.indexer.fused_wqa_wkv" + assert _is_sm70_fp8_prescaled_m1_decode_layer(layer) + layer.prefix = "model.layers.2.attn.wq_b" + assert not _is_sm70_fp8_prescaled_m1_decode_layer(layer) + + +def test_fp8_prescaled_m1_decode_scale_gate_is_reversible(): + scales = torch.tensor([2**-12, 2**-11, 0.0], dtype=torch.float16) + prescaled = _try_sm70_fp8_prescaled_decode_scales(scales) + assert prescaled is not None + torch.testing.assert_close(prescaled, scales * 256, rtol=0, atol=0) + + assert ( + _try_sm70_fp8_prescaled_decode_scales( + torch.tensor([512.0], dtype=torch.float16) + ) + is None + ) + assert ( + _try_sm70_fp8_prescaled_decode_scales( + torch.tensor([-(2**-12)], dtype=torch.float16) + ) + is None + ) + + +def test_fp8_prescaled_m1_decode_only_handles_m1(monkeypatch): + calls = [] + + def fake_default(out, input, qweight, scales, *args): + calls.append(("default", input.shape[0], scales)) + out.zero_() + + def fake_prescaled(out, input, qweight, scales, *args): + calls.append(("prescaled", input.shape[0], scales)) + out.zero_() + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.fp8.sm70_ops.fp8_gemm_sm70_out", + fake_default, + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.fp8.sm70_ops." + "fp8_gemm_sm70_prescaled_m1_out", + fake_prescaled, + ) + normal_scales = torch.empty((1, 6), dtype=torch.float16) + prescaled_scales = torch.empty((1, 6), dtype=torch.float16) + layer = SimpleNamespace( + sm70_fp8_turbomind=True, + sm70_fp8_bmm=False, + output_size_per_partition=6, + weight=torch.empty((4, 6), dtype=torch.uint8), + weight_scale_inv=normal_scales, + sm70_fp8_decode_prescaled_scales=prescaled_scales, + sm70_fp8_k_ld=4, + sm70_fp8_q_ld=6, + ) + method = SimpleNamespace() + + monkeypatch.setenv("VLLM_SM70_FP8_PRESCALED_M1_DECODE", "1") + envs.disable_envs_cache() + try: + Fp8LinearMethod.apply(method, layer, torch.empty((1, 4), dtype=torch.float16)) + Fp8LinearMethod.apply(method, layer, torch.empty((2, 4), dtype=torch.float16)) + monkeypatch.setenv("VLLM_SM70_FP8_PRESCALED_M1_DECODE", "0") + envs.disable_envs_cache() + Fp8LinearMethod.apply(method, layer, torch.empty((1, 4), dtype=torch.float16)) + finally: + envs.disable_envs_cache() + + assert [(route, m) for route, m, _ in calls] == [ + ("prescaled", 1), + ("default", 2), + ("default", 1), + ] + assert calls[0][2] is prescaled_scales + assert calls[1][2] is normal_scales + assert calls[2][2] is normal_scales + + def test_fp8_qpn8_dispatches_small_m_and_workspace_fallback(monkeypatch): calls = [] diff --git a/tests/quantization/test_sm70_glm53_modelopt_nvfp4.py b/tests/quantization/test_sm70_glm53_modelopt_nvfp4.py new file mode 100644 index 0000000000..c712b52c3f --- /dev/null +++ b/tests/quantization/test_sm70_glm53_modelopt_nvfp4.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Narrow SM70 admission and shape gates for GLM-5.3 ModelOpt NVFP4.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from vllm.model_executor.layers.quantization import modelopt +from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm +from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config +from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( + ModelOptNvFp4SM70MoEMethod, + validate_nvfp4_sm70_moe_contract, +) + + +def _glm53_moe_contract(**overrides): + values = { + "num_experts": 288, + "experts_per_token": 8, + "hidden_dim": 4096, + "intermediate_size_per_partition": 512, + "tp_size": 4, + "moe_parallel_config": SimpleNamespace(use_all2all_kernels=False), + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_nvfp4_moe_contract_accepts_glm53_tp4(): + validate_nvfp4_sm70_moe_contract(_glm53_moe_contract()) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("num_experts", 256), + ("experts_per_token", 6), + ("hidden_dim", 2048), + ("intermediate_size_per_partition", 256), + ("tp_size", 2), + ], +) +def test_nvfp4_moe_contract_rejects_unvalidated_glm53_shapes(field, value): + with pytest.raises(NotImplementedError): + validate_nvfp4_sm70_moe_contract(_glm53_moe_contract(**{field: value})) + + +def test_nvfp4_moe_contract_rejects_glm53_tp8(): + with pytest.raises(NotImplementedError, match="tensor parallel"): + validate_nvfp4_sm70_moe_contract( + _glm53_moe_contract(tp_size=8, intermediate_size_per_partition=256) + ) + + +def test_pure_nvfp4_glm53_moe_uses_turbomind_w4a16_on_sm70(): + config = ModelOptNvFp4Config( + quant_method="NVFP4", + is_checkpoint_nvfp4_serialized=True, + ) + + class FakeRoutedExperts: + moe_config = _glm53_moe_contract() + + with ( + patch.object(modelopt, "RoutedExperts", FakeRoutedExperts), + patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=True), + patch.object(sm70_tm, "should_use_nvfp4_moe_turbomind", return_value=True), + ): + method = config.get_quant_method( + FakeRoutedExperts(), "model.layers.3.mlp.experts" + ) + + assert isinstance(method, ModelOptNvFp4SM70MoEMethod) + assert method.use_a16 diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index 065b447cc9..5b6406f0f1 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -19,6 +19,9 @@ from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( ModelOptNvFp4SM70MoEMethod, _prepare_compact_slot_groups, + _prepare_single_token_slots, + _single_token_weighted_reduce, + _use_qwen38_qpn_m1_decode, _validate_weight_layout, validate_nvfp4_sm70_moe_contract, ) @@ -57,6 +60,19 @@ def _moe_contract(**overrides): return SimpleNamespace(**values) +def _qwen4_moe_contract(**overrides): + values = { + "num_experts": 512, + "experts_per_token": 10, + "hidden_dim": 2560, + "intermediate_size_per_partition": 160, + "tp_size": 4, + "moe_parallel_config": SimpleNamespace(use_all2all_kernels=False), + } + values.update(overrides) + return SimpleNamespace(**values) + + def test_mixed_min_capability_requires_exact_sm70_and_both_turbomind_routes(): with ( patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=True), @@ -92,12 +108,42 @@ def test_nvfp4_grouped_prefill_defaults_on_and_can_be_disabled(monkeypatch): ("tp_size", 2), ], ) -def test_nvfp4_moe_contract_rejects_non_qwen36_shapes(field, value): +def test_nvfp4_moe_contract_rejects_unvalidated_shapes(field, value): validate_nvfp4_sm70_moe_contract(_moe_contract()) with pytest.raises(NotImplementedError): validate_nvfp4_sm70_moe_contract(_moe_contract(**{field: value})) +def test_nvfp4_moe_contract_accepts_qwen4_exp_tp4(): + validate_nvfp4_sm70_moe_contract(_qwen4_moe_contract()) + + +def test_qwen38_qpn_m1_decode_is_default_on_and_exact_shape_only(monkeypatch): + layer = SimpleNamespace( + moe_config=_qwen4_moe_contract(), + sm70_nvfp4_num_experts=512, + sm70_nvfp4_hidden_size=2560, + sm70_nvfp4_intermediate_size=160, + sm70_nvfp4_top_k=10, + ) + x = torch.empty(1, 2560, dtype=torch.float16) + topk_ids = torch.empty(1, 10, dtype=torch.int32) + + monkeypatch.delenv("VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE", raising=False) + assert _use_qwen38_qpn_m1_decode(layer, x, topk_ids) + + monkeypatch.setenv("VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE", "0") + assert not _use_qwen38_qpn_m1_decode(layer, x, topk_ids) + + monkeypatch.setenv("VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE", "1") + assert _use_qwen38_qpn_m1_decode(layer, x, topk_ids) + assert not _use_qwen38_qpn_m1_decode(layer, x.repeat(2, 1), topk_ids) + assert not _use_qwen38_qpn_m1_decode(layer, x, topk_ids.long()) + + layer.moe_config.tp_size = 2 + assert not _use_qwen38_qpn_m1_decode(layer, x, topk_ids) + + def test_nvfp4_moe_contract_rejects_shape_consistent_unvalidated_tp8(): with pytest.raises(NotImplementedError, match="tensor parallel"): validate_nvfp4_sm70_moe_contract( @@ -160,7 +206,7 @@ def test_nvfp4_sm70_moe_owns_routing_without_generic_modular_wrapper(): not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), reason="requires an exact SM70 CUDA device", ) -@pytest.mark.parametrize("total_slots", (8, 72, 80)) +@pytest.mark.parametrize("total_slots", (8, 72, 80, 100)) def test_nvfp4_compact_groups_keep_duplicate_expert_slots_independent(total_slots): sorted_expert_ids = ( torch.arange(total_slots, dtype=torch.int32, device="cuda") // 3 @@ -177,6 +223,56 @@ def test_nvfp4_compact_groups_keep_duplicate_expert_slots_independent(total_slot assert torch.equal(active_expert_ids.cpu(), sorted_expert_ids.cpu()) +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), + reason="requires an exact SM70 CUDA device", +) +def test_nvfp4_single_token_direct_routing_is_exact_and_graph_dynamic(): + top_k = 10 + hidden = 2560 + x = torch.randn(1, hidden, dtype=torch.float16, device="cuda") + topk_ids = torch.tensor( + [[401, 7, 310, 99, 211, 3, 470, 120, 55, 256]], + dtype=torch.int64, + device="cuda", + ) + weights = torch.softmax(torch.randn(1, top_k, device="cuda"), dim=-1) + expert_output = torch.randn(top_k, hidden, dtype=torch.float16, device="cuda") + expanded = torch.empty_like(expert_output) + active_ids = torch.empty(top_k, dtype=torch.int32, device="cuda") + output = torch.empty(1, hidden, dtype=torch.float16, device="cuda") + reference = torch.empty_like(output) + identity = torch.arange(top_k, dtype=torch.int32, device="cuda").view(1, -1) + + _prepare_single_token_slots(x, topk_ids, expanded, active_ids) + _single_token_weighted_reduce(expert_output, weights, output) + torch.ops._moe_C.moe_unpermute( + expert_output, weights, identity, None, top_k, reference + ) + torch.accelerator.synchronize() + assert torch.equal(expanded, x.expand(top_k, -1)) + assert torch.equal(active_ids, topk_ids[0].int()) + assert torch.equal(output, reference) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _prepare_single_token_slots(x, topk_ids, expanded, active_ids) + _single_token_weighted_reduce(expert_output, weights, output) + + x.add_(0.25) + topk_ids.copy_(topk_ids.roll(1, dims=1)) + weights.copy_(weights.roll(2, dims=1)) + expert_output.mul_(0.5) + graph.replay() + torch.ops._moe_C.moe_unpermute( + expert_output, weights, identity, None, top_k, reference + ) + torch.accelerator.synchronize() + assert torch.equal(expanded, x.expand(top_k, -1)) + assert torch.equal(active_ids, topk_ids[0].int()) + assert torch.equal(output, reference) + + def test_mixed_w4a16_moe_requires_turbomind_on_sm70(): config = _mixed_config() @@ -190,3 +286,25 @@ class FakeRoutedExperts: pytest.raises(NotImplementedError, match="TurboMind"), ): config.get_quant_method(FakeRoutedExperts(), "model.layers.0.mlp.experts") + + +def test_pure_nvfp4_qwen4_moe_uses_turbomind_w4a16_on_sm70(): + config = ModelOptNvFp4Config( + quant_method="NVFP4", + is_checkpoint_nvfp4_serialized=True, + ) + + class FakeRoutedExperts: + moe_config = _qwen4_moe_contract() + + with ( + patch.object(modelopt, "RoutedExperts", FakeRoutedExperts), + patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=True), + patch.object(sm70_tm, "should_use_nvfp4_moe_turbomind", return_value=True), + ): + method = config.get_quant_method( + FakeRoutedExperts(), "model.layers.0.mlp.experts" + ) + + assert isinstance(method, ModelOptNvFp4SM70MoEMethod) + assert method.use_a16 diff --git a/tests/quantization/test_sm70_mxfp4_moe.py b/tests/quantization/test_sm70_mxfp4_moe.py index 1f2e7076e4..2a28c2c008 100644 --- a/tests/quantization/test_sm70_mxfp4_moe.py +++ b/tests/quantization/test_sm70_mxfp4_moe.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.quantization.mxfp4_sm70_moe import ( Mxfp4SM70MoEMethod, _compact_mxfp4_active_experts, + _mxfp4_qpn_m1_decode_contract, _select_mxfp4_direct_order_offsets, _select_mxfp4_stage_dispatch, validate_mxfp4_sm70_moe_contract, @@ -213,8 +214,7 @@ def test_mxfp4_sm70_b1_dispatch_selects_six_runtime_experts(monkeypatch): def test_mxfp4_sm70_direct_order_ignores_compacted_offsets(): buffers = { - # Reproduce the state left by an M8 active-expert compaction before - # the B1 graph is captured. + # Reproduce M8 compaction state left before the B1 graph is captured. "compact_expert_offsets": torch.tensor( [0, 8, 9, 16, 24, 32, 48], dtype=torch.int32 ), @@ -227,6 +227,61 @@ def test_mxfp4_sm70_direct_order_ignores_compacted_offsets(): torch.testing.assert_close(offsets, torch.arange(7, dtype=torch.int32)) +def test_mxfp4_sm70_qpn_m1_is_default_with_rollback(monkeypatch): + monkeypatch.delenv("VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE", raising=False) + envs.disable_envs_cache() + try: + assert envs.VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE + + monkeypatch.setenv("VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE", "0") + envs.disable_envs_cache() + assert not envs.VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE + finally: + envs.disable_envs_cache() + + +def test_mxfp4_sm70_qpn_m1_missing_default_op_falls_back(monkeypatch): + monkeypatch.delenv("VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE", raising=False) + monkeypatch.setattr(mxfp4_moe, "_mxfp4_qpn_m1_op_available", lambda: False) + envs.disable_envs_cache() + try: + assert not mxfp4_moe._mxfp4_qpn_m1_extension_enabled() + + monkeypatch.setenv("VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE", "1") + envs.disable_envs_cache() + with pytest.raises(RuntimeError, match="explicitly enabled"): + mxfp4_moe._mxfp4_qpn_m1_extension_enabled() + finally: + envs.disable_envs_cache() + + +def test_mxfp4_sm70_qpn_m1_requires_exact_tp4_tensor_contract(monkeypatch): + monkeypatch.setattr(envs, "VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE", True) + layer = SimpleNamespace( + moe_config=SimpleNamespace(tp_size=4), + local_num_experts=256, + global_num_experts=256, + sm70_mxfp4_qpn_m1_available=True, + sm70_mxfp4_w13_k_dim=4096, + sm70_mxfp4_w13_n_dim=1024, + sm70_mxfp4_w2_k_dim=512, + sm70_mxfp4_w2_n_dim=4096, + sm70_mxfp4_group_size=32, + w13_tm_weight=torch.empty((256, 4096, 128), device="meta"), + w13_tm_scales=torch.empty((256, 128, 1024), device="meta"), + w2_tm_weight=torch.empty((256, 512, 512), device="meta"), + w2_tm_scales=torch.empty((256, 16, 4096), device="meta"), + ) + + assert _mxfp4_qpn_m1_decode_contract(layer, direct_order=True) + assert not _mxfp4_qpn_m1_decode_contract(layer, direct_order=False) + layer.sm70_mxfp4_qpn_m1_available = False + assert not _mxfp4_qpn_m1_decode_contract(layer, direct_order=True) + layer.sm70_mxfp4_qpn_m1_available = True + layer.moe_config.tp_size = 8 + assert not _mxfp4_qpn_m1_decode_contract(layer, direct_order=True) + + def test_mxfp4_sm70_b1_dispatch_rejects_incompatible_permute_fastpath( monkeypatch, ): diff --git a/tests/quantization/test_sm70_nvfp4_qpn2.py b/tests/quantization/test_sm70_nvfp4_qpn2.py index 84f8174158..c9084e7cd6 100644 --- a/tests/quantization/test_sm70_nvfp4_qpn2.py +++ b/tests/quantization/test_sm70_nvfp4_qpn2.py @@ -18,12 +18,20 @@ def test_nvfp4_qpn2_is_default_off_with_explicit_on(monkeypatch): monkeypatch.delenv("VLLM_SM70_NVFP4_QPN2", raising=False) + monkeypatch.delenv("VLLM_SM70_NVFP4_QPN2_PREFILL", raising=False) + monkeypatch.delenv("VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M", raising=False) envs.disable_envs_cache() try: assert not envs.VLLM_SM70_NVFP4_QPN2 + assert not envs.VLLM_SM70_NVFP4_QPN2_PREFILL + assert envs.VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M == 1024 monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2", "1") + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2_PREFILL", "1") + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M", "9") envs.disable_envs_cache() assert envs.VLLM_SM70_NVFP4_QPN2 + assert envs.VLLM_SM70_NVFP4_QPN2_PREFILL + assert envs.VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M == 9 finally: envs.disable_envs_cache() @@ -75,6 +83,8 @@ def _make_small_layer() -> torch.nn.Module: def test_nvfp4_qpn2_prepare_and_dispatch_contract(monkeypatch): monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2", "1") + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2_PREFILL", "1") + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M", "9") envs.disable_envs_cache() layer = _make_small_layer() calls = [] @@ -88,6 +98,13 @@ def test_nvfp4_qpn2_prepare_and_dispatch_contract(monkeypatch): ) monkeypatch.setattr(nvfp4_scheme, "_is_qpn2_layer", lambda layer: True) monkeypatch.setattr(nvfp4_scheme, "_missing_qpn2_ops", lambda: []) + monkeypatch.setattr(nvfp4_scheme, "_missing_qpn2_prefill_ops", lambda: []) + workspace = torch.empty((1,), dtype=torch.float16) + monkeypatch.setattr( + nvfp4_scheme.sm70_tm, + "get_nvfp4_qpn4_dense_workspace", + lambda _weight: workspace, + ) monkeypatch.setitem(nvfp4_scheme._SM70_NVFP4_QPN2_CONFIGS, (64, 64, False), (8, 2)) monkeypatch.setitem(nvfp4_scheme._SM70_NVFP4_QPN2_CONFIGS, (64, 64, True), (8, 2)) monkeypatch.setattr( @@ -121,12 +138,34 @@ def fake_dispatch(*args): monkeypatch.setattr( nvfp4_scheme.sm70_ops, "nvfp4_qpn2_dispatch_sm70_out", fake_dispatch ) + prefill_calls = [] + + def fake_prefill(*args): + prefill_calls.append(args) + args[0].fill_(5) + + monkeypatch.setattr( + nvfp4_scheme.sm70_ops, + "nvfp4_qpn4_prefill_sm70_out", + fake_prefill, + ) try: scheme.process_weights_after_loading(layer) assert layer.sm70_nvfp4_qpn2 assert layer.sm70_nvfp4_qpn2_gated_silu assert layer.sm70_nvfp4_qpn2_global_scale == 0.5 + assert layer.sm70_nvfp4_qpn2_prefill_dense_weight_ptr == workspace.data_ptr() + assert layer.sm70_nvfp4_qpn2_prefill_codes.shape == (64, 32) + assert layer.sm70_nvfp4_qpn2_prefill_scales.shape == (4, 64) + assert ( + layer.sm70_nvfp4_qpn2_prefill_codes.data_ptr() + == layer.sm70_nvfp4_qpn2_codes.data_ptr() + ) + assert ( + layer.sm70_nvfp4_qpn2_prefill_scales.data_ptr() + == layer.sm70_nvfp4_qpn2_scales.data_ptr() + ) assert layer.weight.numel() == 0 assert layer.weight_scale.numel() == 0 @@ -139,5 +178,16 @@ def fake_dispatch(*args): assert torch.equal(fused, torch.full_like(fused, 3)) assert calls[0][-1] is False assert calls[1][-1] is True + + large_x = torch.ones((9, 64), dtype=torch.float16) + large_raw = scheme.apply_weights(layer, large_x) + large_fused = scheme.apply_fused_silu_and_mul(layer, large_x) + assert torch.equal(large_raw, torch.full_like(large_raw, 5)) + assert large_fused is not None + assert torch.equal(large_fused, torch.full_like(large_fused, 5)) + assert prefill_calls[0][-2:] == (True, False) + assert prefill_calls[1][-2:] == (True, True) + assert prefill_calls[0][3].shape == (64, 32) + assert prefill_calls[0][4].shape == (4, 64) finally: envs.disable_envs_cache() diff --git a/tests/quantization/test_sm70_online_qpn8.py b/tests/quantization/test_sm70_online_qpn8.py new file mode 100644 index 0000000000..cd2aa75285 --- /dev/null +++ b/tests/quantization/test_sm70_online_qpn8.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.envs as envs +from vllm.model_executor.layers.quantization import sm70_online_qpn8 as online_qpn8 + + +def test_qwen4_exp_online_qpn8_defaults_on_and_can_be_disabled(monkeypatch): + monkeypatch.delenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8", raising=False) + envs.disable_envs_cache() + try: + assert envs.VLLM_SM70_QWEN4_EXP_ONLINE_QPN8 + monkeypatch.setenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8", "0") + envs.disable_envs_cache() + assert not envs.VLLM_SM70_QWEN4_EXP_ONLINE_QPN8 + finally: + envs.disable_envs_cache() + + +def test_qwen3next_shared_gate_fusion_defaults_on_and_can_be_disabled(monkeypatch): + monkeypatch.delenv("VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION", raising=False) + envs.disable_envs_cache() + try: + assert envs.VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION + monkeypatch.setenv("VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION", "0") + envs.disable_envs_cache() + assert not envs.VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION + finally: + envs.disable_envs_cache() + + +def test_default_on_qpn_sidecars_load_without_enable_overrides(monkeypatch): + calls: list[str] = [] + monkeypatch.setenv("VLLM_SM70_FP8_QPN8_LIBRARY", "/tmp/qpn8.so") + monkeypatch.setenv("VLLM_SM70_NVFP4_QPN_M1_LIBRARY", "/tmp/qpn-m1.so") + monkeypatch.delenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8", raising=False) + monkeypatch.delenv("VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE", raising=False) + monkeypatch.setattr(torch.ops, "load_library", calls.append) + + online_qpn8.sm70_ops._maybe_load_fp8_qpn8_library() + online_qpn8.sm70_ops._maybe_load_nvfp4_qpn_m1_library() + + assert calls == ["/tmp/qpn8.so", "/tmp/qpn-m1.so"] + + +@pytest.mark.parametrize( + ("prefix", "k", "n", "expected"), + [ + ("model.layers.0.linear_attn.in_proj_qkvz", 2560, 4096, (16, 1, False)), + ("model.layers.0.linear_attn.out_proj", 1536, 2560, (12, 2, False)), + ("model.layers.3.self_attn.qkv_proj", 2560, 3584, (16, 1, False)), + ("model.layers.3.self_attn.o_proj", 1536, 2560, (12, 2, False)), + ( + "model.layers.0.attn_hyper_connection.input_mix_weight_up", + 320, + 10240, + (4, 2, False), + ), + ( + "model.layers.0.attn_hyper_connection.input_mix_weight_down_block_inject", + 10240, + 336, + (32, 1, False), + ), + ("model.layers.0.mlp.shared_expert.gate_up_proj", 2560, 320, None), + ], +) +def test_online_qpn8_shape_gate(prefix, k, n, expected): + assert online_qpn8._shape_config(prefix, k, n) == expected + + +def test_online_qpn8_runtime_contract_is_tp4_no_mtp(monkeypatch): + text_config = SimpleNamespace( + model_type="generic_hybrid_moe", + hidden_size=2560, + num_hidden_layers=48, + num_experts=512, + num_experts_per_tok=10, + moe_intermediate_size=640, + hc_count=4, + hc_lowrank=320, + num_attention_heads=24, + num_key_value_heads=2, + indexer_head_dim=128, + indexer_budget=2048, + indexer_compress_ratio=4, + ) + config = SimpleNamespace( + model_config=SimpleNamespace(hf_text_config=text_config), + speculative_config=None, + ) + monkeypatch.setattr(online_qpn8, "get_current_vllm_config", lambda: config) + monkeypatch.setattr(online_qpn8, "get_tensor_model_parallel_world_size", lambda: 4) + assert online_qpn8._exact_runtime_contract() + + config.speculative_config = SimpleNamespace(method="mtp") + assert not online_qpn8._exact_runtime_contract() + config.speculative_config = None + monkeypatch.setattr(online_qpn8, "get_tensor_model_parallel_world_size", lambda: 2) + assert not online_qpn8._exact_runtime_contract() + monkeypatch.setattr(online_qpn8, "get_tensor_model_parallel_world_size", lambda: 4) + text_config.hc_count = 2 + assert not online_qpn8._exact_runtime_contract() + + +def test_online_qpn8_apply_uses_prepared_state(monkeypatch): + layer = torch.nn.Module() + layer.register_buffer( + online_qpn8._CODES_ATTR, torch.empty((16, 32), dtype=torch.uint8) + ) + layer.register_buffer( + online_qpn8._SCALES_ATTR, torch.empty((1, 32), dtype=torch.float16) + ) + layer._sm70_qwen4_exp_online_qpn8 = True + layer._sm70_qwen4_exp_online_qpn8_k = 16 + layer._sm70_qwen4_exp_online_qpn8_n = 32 + layer._sm70_qwen4_exp_online_qpn8_logical_n = 32 + layer._sm70_qwen4_exp_online_qpn8_split_k = 4 + layer._sm70_qwen4_exp_online_qpn8_nacc = 2 + layer._sm70_qwen4_exp_online_qpn8_prefetch = False + layer._sm70_qwen4_exp_online_qpn8_workspace_ptr = 123 + calls = [] + + def fake_dispatch(*args): + calls.append(args) + args[0].fill_(2) + + monkeypatch.setattr( + online_qpn8.sm70_ops, "fp8_qpn8_dispatch_sm70_out", fake_dispatch + ) + x = torch.ones((2, 3, 16), dtype=torch.float16) + bias = torch.ones((32,), dtype=torch.float16) + out = online_qpn8.maybe_apply_online_qpn8(layer, x, bias) + assert out is not None and out.shape == (2, 3, 32) + assert torch.equal(out, torch.full_like(out, 3)) + assert calls[0][1] == 123 + assert calls[0][-1] is False + + +def test_fused_hc_uses_prepared_partials_without_global_lookup(monkeypatch): + down = torch.nn.Module() + up = torch.nn.Module() + for layer, k, n in ((down, 10240, 384), (up, 320, 10240)): + setattr(layer, online_qpn8._STATE_ATTR, True) + layer._sm70_qwen4_exp_online_qpn8_k = k + layer._sm70_qwen4_exp_online_qpn8_n = n + layer._sm70_qwen4_exp_online_qpn8_workspace_ptr = 123 + layer.register_buffer( + online_qpn8._CODES_ATTR, torch.empty((1,), dtype=torch.uint8) + ) + layer.register_buffer( + online_qpn8._SCALES_ATTR, torch.empty((1,), dtype=torch.float16) + ) + partials = torch.empty((32 * 384,), dtype=torch.float32) + down.register_buffer( + online_qpn8._HC_PARTIALS_ATTR, + partials, + persistent=False, + ) + calls = [] + + def fake_dispatch(*args): + calls.append(args) + args[0].fill_(2) + args[1].fill_(3) + + monkeypatch.setattr(online_qpn8.sm70_ops, "has_fp8_qpn8_hc_dispatch", lambda: True) + monkeypatch.setattr( + online_qpn8.sm70_ops, "fp8_qpn8_hc_dispatch_sm70_out", fake_dispatch + ) + xn = torch.ones((1, 10240), dtype=torch.float16) + result = online_qpn8.maybe_apply_fused_hc(down, up, xn) + assert result is not None + block_out, injection_out = result + assert torch.equal(block_out, torch.full_like(block_out, 2)) + assert torch.equal(injection_out, torch.full_like(injection_out, 3)) + assert calls[0][5] is partials + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_channel_qpn8_accepts_k16_n32_alignment(): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + n, k = 64, 320 + weight = torch.randn((n, k), dtype=torch.float16, device="cuda").mul_(0.02) + scales = weight.float().abs().amax(dim=1, keepdim=True).div_(448.0) + qweight = (weight.float() / scales).to(torch.float8_e4m3fn) + try: + codes, packed_scales = online_qpn8.sm70_ops.fp8_qpn8_prepare_sm70( + qweight.contiguous(), scales.contiguous() + ) + except RuntimeError as exc: + if "multiples of 128" in str(exc): + pytest.skip("installed SM70 extension predates channel-QPN8 alignment") + raise + x = torch.randn((1, k), dtype=torch.float16, device="cuda").mul_(0.1) + out = torch.empty((1, n), dtype=torch.float16, device="cuda") + online_qpn8.sm70_ops.fp8_qpn8_gemm_sm70_out( + out, x, codes, packed_scales, 4, 2, True, False + ) + reference = x.float().matmul((qweight.float() * scales).t()).half() + torch.accelerator.synchronize() + relative_l2 = (out.float() - reference.float()).norm() / reference.float().norm() + assert float(relative_l2) < 5e-3 diff --git a/tests/test_envs.py b/tests/test_envs.py index 67a8f41484..983359d14f 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -147,12 +147,28 @@ def test_sm70_concurrency_tuning_envs( monkeypatch.setenv(name, "0") assert environment_variables[name]() is False + name = "VLLM_SM70_FP8_PRESCALED_M1_DECODE" + monkeypatch.delenv(name, raising=False) + assert environment_variables[name]() is True + monkeypatch.setenv(name, "0") + assert environment_variables[name]() is False + name = "VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE" monkeypatch.delenv(name, raising=False) assert environment_variables[name]() is True monkeypatch.setenv(name, "0") assert environment_variables[name]() is False + for name in ( + "VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE", + "VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE", + "VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE", + ): + monkeypatch.delenv(name, raising=False) + assert environment_variables[name]() is True + monkeypatch.setenv(name, "0") + assert environment_variables[name]() is False + monkeypatch.delenv("VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS", raising=False) assert environment_variables["VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS"]() == 128 monkeypatch.setenv("VLLM_SM70_NVFP4_MOE_TUNE_MAX_TOKENS", "640") @@ -202,6 +218,19 @@ def test_flash_v100_e4m3_page800_fastpath_envs( assert environment_variables[name]() is not expected_bool +def test_flash_v100_dflash2_grouped_layout_envs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + for name in ( + "VLLM_FLASH_V100_DFLASH2_FIXED_INTERLEAVED", + "VLLM_FLASH_V100_DFLASH2_STAGE_PAGE_IDS", + ): + monkeypatch.delenv(name, raising=False) + assert environment_variables[name]() is True + monkeypatch.setenv(name, "0") + assert environment_variables[name]() is False + + class TestEnvWithChoices: """Test cases for env_with_choices function.""" diff --git a/tests/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py index ca2a60d4ce..dbeeb61e1d 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -482,6 +482,162 @@ def test_extract_tool_calls_fallback_no_tags( assert extracted_tool_calls.tool_calls[0].function.name == "get_current_weather" +def test_extract_json_body_tool_call(qwen3_tool_parser): + """Accept the alternate JSON body emitted by Qwen/OpenCode prompts.""" + model_output = ( + "I will check it." + "\n" + '{"name":"get_current_weather","arguments":' + '{"city":"Dallas","state":"TX"}}\n' + "" + ) + request = ChatCompletionRequest(model=MODEL, messages=[]) + + extracted = qwen3_tool_parser.extract_tool_calls(model_output, request=request) + + assert extracted.tools_called + assert extracted.content == "I will check it." + assert len(extracted.tool_calls) == 1 + assert extracted.tool_calls[0].function.name == "get_current_weather" + assert json.loads(extracted.tool_calls[0].function.arguments) == { + "city": "Dallas", + "state": "TX", + } + + +def test_streaming_json_body_tool_call(qwen3_tool_parser): + """Buffer split JSON arguments and emit one valid atomic tool delta.""" + request = ChatCompletionRequest(model=MODEL, messages=[]) + deltas = [ + "I will check it.\n{", + '"name":"get_current_weather",', + '"arguments":{"city":"Dallas",', + '"state":"TX"}}\n', + ] + + from tests.tool_parsers.utils import ( + run_tool_extraction_streaming, + ) + + reconstructor = run_tool_extraction_streaming( + qwen3_tool_parser, + deltas, + request, + assert_one_tool_per_delta=False, + ) + + assert reconstructor.other_content == "I will check it." + assert len(reconstructor.tool_calls) == 1 + assert reconstructor.tool_calls[0].function.name == "get_current_weather" + assert json.loads(reconstructor.tool_calls[0].function.arguments) == { + "city": "Dallas", + "state": "TX", + } + + +def test_malformed_json_body_remains_content(qwen3_tool_parser): + model_output = '{"name":"get_current_weather","arguments":{' + request = ChatCompletionRequest(model=MODEL, messages=[]) + + extracted = qwen3_tool_parser.extract_tool_calls(model_output, request=request) + + assert not extracted.tools_called + assert extracted.content == model_output + + +def test_parameter_markers_inside_string_value(qwen3_tokenizer): + """Write/Edit payloads may contain XML-like text without adding fields.""" + tools = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "Write", + "parameters": { + "type": "object", + "properties": { + "file_path": {"type": "string"}, + "content": {"type": "string"}, + }, + "required": ["file_path", "content"], + "additionalProperties": False, + }, + }, + ) + ] + content = '{"open":"literal","close":"beforeafter"}' + model_output = ( + "\n" + "\n" + "\n/tmp/result.json\n\n" + f"\n{content}\n\n" + "\n" + "" + ) + request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) + + parser = Qwen3CoderToolParser(qwen3_tokenizer, tools=tools) + extracted = parser.extract_tool_calls(model_output, request=request) + + assert extracted.tools_called + assert len(extracted.tool_calls) == 1 + assert json.loads(extracted.tool_calls[0].function.arguments) == { + "file_path": "/tmp/result.json", + "content": content, + } + + +def test_streaming_parameter_markers_inside_string_value(qwen3_tokenizer): + """Wait for structural lookahead when a chunk ends on a literal close tag.""" + tools = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "Write", + "parameters": { + "type": "object", + "properties": { + "file_path": {"type": "string"}, + "content": {"type": "string"}, + }, + "required": ["file_path", "content"], + "additionalProperties": False, + }, + }, + ) + ] + content = '{"open":"literal","close":"beforeafter"}' + deltas = [ + "", + "\n", + "\n", + "\n/tmp/result.json\n\n", + ( + "\n" + '{"open":"literal","close":"before' + ), + 'after"}\n\n\n', + ] + request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) + + from tests.tool_parsers.utils import ( + run_tool_extraction_streaming, + ) + + parser = Qwen3CoderToolParser(qwen3_tokenizer, tools=tools) + reconstructor = run_tool_extraction_streaming( + parser, + deltas, + request, + assert_one_tool_per_delta=False, + ) + + assert len(reconstructor.tool_calls) == 1 + assert json.loads(reconstructor.tool_calls[0].function.arguments) == { + "file_path": "/tmp/result.json", + "content": content, + } + + def test_extract_tool_calls_type_conversion(qwen3_tokenizer): """Test parameter type conversion based on tool schema""" tools = [ @@ -1381,10 +1537,98 @@ def test_streaming_multi_param_single_chunk(qwen3_tool_parser, qwen3_tokenizer): assert args["unit"] == "fahrenheit" +def test_streaming_final_param_and_function_close_single_chunk( + qwen3_tool_parser, +): + """A speculative burst must not omit the final JSON closing brace.""" + request = ChatCompletionRequest(model=MODEL, messages=[]) + deltas = [ + "", + "\n", + "\n", + ( + "\nDallas\n\n" + "\nTX\n\n" + "\n" + ), + ] + + from tests.tool_parsers.utils import ( + run_tool_extraction_streaming, + ) + + reconstructor = run_tool_extraction_streaming( + qwen3_tool_parser, + deltas, + request, + assert_one_tool_per_delta=False, + ) + + assert len(reconstructor.tool_calls) == 1 + assert json.loads(reconstructor.tool_calls[0].function.arguments) == { + "city": "Dallas", + "state": "TX", + } + + +def test_streaming_parallel_function_body_buffered_before_header_emit( + qwen3_tokenizer, +): + """A final speculative burst must not leave the last call as ``{}``.""" + number_params = { + "type": "object", + "properties": {"number": {"type": "integer"}}, + "required": ["number"], + } + tools = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "math.factorial", + "description": "Calculate a factorial.", + "parameters": number_params, + }, + ) + ] + parser = Qwen3CoderToolParser(qwen3_tokenizer, tools=tools) + request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) + # These are the observed DFlash2 q15 streaming boundaries. The second + # function closes in the same burst that starts the third one; by the time + # the third header is emitted, its complete body is already buffered. + deltas = [ + "", + "\n", + "\n\n5", + "\n\n", + "\n\n", + "\n\n\n", + "1", + "0\n\n\n\n\n\n\n15", + "\n\n", + "\n", + ] + + from tests.tool_parsers.utils import run_tool_extraction_streaming + + reconstructor = run_tool_extraction_streaming( + parser, + deltas, + request, + assert_one_tool_per_delta=False, + ) + + assert [ + json.loads(call.function.arguments) for call in reconstructor.tool_calls + ] == [{"number": 5}, {"number": 10}, {"number": 15}] + + @pytest.mark.parametrize( "deltas, expected_content", [ - (["<", "meta charset=\"UTF-8\">"], ''), + (["<", 'meta charset="UTF-8">'], ''), (["macOS"], "macOS"), (["<", "<", "span>text"], "<text"), ], diff --git a/tests/transformers_utils/processors/test_glm5next.py b/tests/transformers_utils/processors/test_glm5next.py new file mode 100644 index 0000000000..12b6239e06 --- /dev/null +++ b/tests/transformers_utils/processors/test_glm5next.py @@ -0,0 +1,401 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the vLLM-native GLM-5.3-Flash multimodal processor.""" + +import math + +import pytest +import torch +from PIL import Image +from transformers.image_utils import PILImageResampling +from transformers.video_utils import VideoMetadata + +from vllm.transformers_utils.processors.glm5next import ( + Glm5NextImageProcessor, + Glm5NextProcessor, + Glm5NextVideoProcessor, + _get_pad_content_size, + _resize_or_pad, + glm_sample_frame_indices, + smart_resize, +) + +PATCH_SIZE = 14 +MERGE_SIZE = 2 +PATCH_EXPAND_FACTOR = 1 # checkpoint processor_config.json +FACTOR = PATCH_SIZE * MERGE_SIZE * PATCH_EXPAND_FACTOR # 28 +PIXELS_PER_TOKEN = 2 * (PATCH_SIZE * MERGE_SIZE) ** 2 # 1568 +MIN_TOKENS = 16 +MAX_TOKENS = 800 # serving-cap equivalent budget, keeps test canvases small +MIN_PIXELS = MIN_TOKENS * PIXELS_PER_TOKEN +MAX_PIXELS = MAX_TOKENS * PIXELS_PER_TOKEN + + +def resize(height: int, width: int, t: int = 2) -> tuple[int, int]: + return smart_resize( + t, + height, + width, + t_factor=2, + h_factor=FACTOR, + w_factor=FACTOR, + min_pixels=MIN_PIXELS, + max_pixels=MAX_PIXELS, + ) + + +@pytest.mark.parametrize( + ("height", "width", "t", "expected"), + [ + (300, 400, 2, (308, 420)), # ceil alignment, no budget pressure + (50, 80, 2, (112, 168)), # below min budget: proportional upscale + (57, 3, 2, (504, 28)), # slim side ceils to one factor + (2160, 3840, 2, (588, 1064)), # shrink: max canvas under the cap + (112, 11200, 2, (84, 7420)), # extreme aspect stays proportional + (300, 400, 16, (252, 308)), # frame count eats into the budget + (112, 11200, 16, (28, 2800)), # slim side stays positive, budget holds + ], +) +def test_smart_resize_reference_values(height, width, t, expected): + got = smart_resize( + t, + height, + width, + t_factor=2, + h_factor=FACTOR, + w_factor=FACTOR, + min_pixels=MIN_PIXELS, + max_pixels=MAX_PIXELS, + ) + assert got == expected + # The aligned canvas never exceeds the budget it was fitted against. + t_bar = max(2, round(t / 2) * 2) + assert t_bar * got[0] * got[1] <= MAX_PIXELS + + +def test_smart_resize_stays_snapped_and_positive(): + heights = (1, 27, 57, 111, 113, 300, 720, 2160) + widths = (1, 27, 57, 111, 113, 400, 1280, 11200) + for h in heights: + for w in widths: + for t in (2, 4, 16, 64): + resized_h, resized_w = resize(h, w, t) + assert resized_h > 0 and resized_w > 0 + assert resized_h % FACTOR == 0 + assert resized_w % FACTOR == 0 + t_bar = max(2, round(t / 2) * 2) + assert t_bar * resized_h * resized_w <= MAX_PIXELS + + +@pytest.mark.parametrize( + ("height", "width", "match"), + [ + (0, 100, "must be positive"), + (100, 0, "must be positive"), + ], +) +def test_smart_resize_rejects_degenerate_inputs(height, width, match): + with pytest.raises(ValueError, match=match): + resize(height, width, 2) + + +def test_smart_resize_rejects_inverted_budget(): + with pytest.raises(ValueError, match="min_pixels must be less than or equal"): + smart_resize( + 2, + 100, + 100, + t_factor=2, + h_factor=FACTOR, + w_factor=FACTOR, + min_pixels=MAX_PIXELS, + max_pixels=MIN_PIXELS, + ) + + +def test_get_pad_content_size(): + # Oversized content shrinks proportionally, never upscales by default. + assert _get_pad_content_size(300, 400, 308, 420) == (300, 400) + assert _get_pad_content_size(600, 800, 308, 420) == (308, 410) + # allow_upscale enlarges small content toward the canvas. + assert _get_pad_content_size(28, 28, 112, 112, allow_upscale=True) == (112, 112) + + +def test_resize_or_pad_pads_right_and_bottom(): + stacked = torch.rand(1, 3, 300, 400) + + def identity_resize(x, size, resample=None): + assert (size.height, size.width) == (300, 400) + return x + + padded = _resize_or_pad(stacked, 308, 420, "pad", None, identity_resize) + assert padded.shape == (1, 3, 308, 420) + torch.testing.assert_close(padded[..., :300, :400], stacked) + assert padded[..., 300:, :].eq(0).all() + assert padded[..., :, 400:].eq(0).all() + + def force_resize(x, size, resample=None): + return torch.zeros(x.shape[0], x.shape[1], size.height, size.width) + + assert _resize_or_pad(stacked, 308, 420, "resize", None, force_resize).shape == ( + 1, + 3, + 308, + 420, + ) + with pytest.raises(ValueError, match="resize_mode"): + _resize_or_pad(stacked, 308, 420, "crop", None, identity_resize) + + +@pytest.fixture(scope="module") +def image_processor(): + return Glm5NextImageProcessor( + min_image_tokens=MIN_TOKENS, max_image_tokens=MAX_TOKENS + ) + + +@pytest.mark.parametrize( + ("height", "width"), + [(50, 80), (300, 400), (2160, 3840), (112, 11200)], +) +def test_image_grid_matches_smart_resize(image_processor, height, width): + out = image_processor(Image.new("RGB", (width, height)), return_tensors="pt") + resized_h, resized_w = resize(height, width) + grid = out["image_grid_thw"][0].tolist() + assert grid == [1, resized_h // PATCH_SIZE, resized_w // PATCH_SIZE] + assert out["pixel_values"].shape == ( + grid[1] * grid[2], + 3 * 2 * PATCH_SIZE * PATCH_SIZE, + ) + + +@pytest.mark.parametrize( + ("height", "width"), + [(50, 80), (300, 400), (2160, 3840), (112, 11200)], +) +def test_image_patch_count_matches_preprocess(image_processor, height, width): + out = image_processor(Image.new("RGB", (width, height)), return_tensors="pt") + grid = out["image_grid_thw"][0].tolist() + expected = grid[1] * grid[2] + assert image_processor.get_number_of_image_patches(height, width) == expected + + +@pytest.fixture(scope="module") +def video_processor(): + return Glm5NextVideoProcessor( + min_image_tokens=MIN_TOKENS, max_image_tokens=MAX_TOKENS + ) + + +def run_video_preprocess(video_processor, frames, **overrides): + kwargs = dict( + resample=PILImageResampling.BICUBIC, + image_mean=(0.48145466, 0.4578275, 0.40821073), + image_std=(0.26862954, 0.26130258, 0.27577711), + patch_size=PATCH_SIZE, + temporal_patch_size=2, + merge_size=MERGE_SIZE, + patch_expand_factor=PATCH_EXPAND_FACTOR, + return_tensors="pt", + ) + kwargs.update(overrides) + return video_processor._preprocess([frames], **kwargs) + + +def test_video_preprocess_pads_odd_frame_count(video_processor): + out = run_video_preprocess(video_processor, torch.rand(7, 3, 300, 400)) + grid = out["video_grid_thw"][0].tolist() + # 7 frames padded to 8 -> grid_t 4; canvas ceil-aligned (308, 420). + assert grid == [4, 22, 30] + assert out["pixel_values_videos"].shape == ( + 4 * 22 * 30, + 3 * 2 * PATCH_SIZE * PATCH_SIZE, + ) + + +def test_video_patch_count_and_video_only_token_count(video_processor): + out = run_video_preprocess(video_processor, torch.rand(7, 3, 300, 400)) + expected_patches = math.prod(out["video_grid_thw"][0].tolist()) + assert video_processor.get_number_of_video_patches(7, 300, 400) == expected_patches + + processor = object.__new__(Glm5NextProcessor) + processor.video_processor = video_processor + token_counts = processor._get_num_multimodal_tokens(video_sizes=[(7, 300, 400)]) + assert token_counts["num_video_tokens"] == [expected_patches // MERGE_SIZE**2] + + +def test_video_preprocess_keeps_min_side_under_frame_budget(video_processor): + out = run_video_preprocess(video_processor, torch.rand(16, 3, 112, 11200)) + # Budget-fitted canvas (28, 2800): 16 * 28 * 2800 == the pixel cap. + assert out["video_grid_thw"][0].tolist() == [8, 2, 200] # height 28, not 0 + + +def test_video_preprocess_pads_content_not_distort(video_processor): + frames = torch.zeros(4, 3, 300, 400) + frames[..., 50, 60] = 1.0 # marker inside the content area + # Bypass rescale/normalize so zero padding stays exactly zero. + out = run_video_preprocess( + video_processor, frames, do_rescale=False, do_normalize=False + ) + # Pad mode keeps the 300x400 content aspect on the (308, 420) canvas + # with zero padding on the right/bottom. + grid = out["video_grid_thw"][0].tolist() + assert grid == [2, 22, 30] + patches = out["pixel_values_videos"].view( + grid[0], + grid[1] // MERGE_SIZE, + grid[2] // MERGE_SIZE, + MERGE_SIZE, + MERGE_SIZE, + 3 * 2 * PATCH_SIZE * PATCH_SIZE, + ) + frame = patches[0].permute(0, 3, 1, 4, 2).reshape(grid[1], grid[2], -1) + # Patch columns fully right of the 400px content are pure padding. + assert frame[:, math.ceil(400 / PATCH_SIZE) :].abs().max() == 0 + + +@pytest.mark.parametrize( + ("total_frames", "fps", "duration", "expected_len", "first", "last"), + [ + # Dense source: the tp-scaled greedy overshoots extract_t, so the + # fixup spreads picks uniformly across the whole video (linspace). + (900, 30.0, 30.0, 60, 0, 899), + (3000, 30.0, 100.0, 200, 0, 2999), + # extract_t capped at 2048 frames. + (72000, 30.0, 2400.0, 2048, 0, 71999), + # Low container fps: the greedy picks every frame. + (48, 2.0, 24.0, 48, 0, 47), + # Duration derived from the frame count when metadata lacks it. + (300, 25.0, 0, 26, 0, 299), + ], +) +def test_glm_sample_frame_indices_behaviour( + total_frames, fps, duration, expected_len, first, last +): + indices = glm_sample_frame_indices(total_frames, fps, duration) + assert len(indices) == expected_len + assert indices[0] == first + assert indices[-1] == last + assert len(indices) % 2 == 0 + assert indices == sorted(indices) + assert all(0 <= idx < total_frames for idx in indices) + + +def test_glm_sample_frame_indices_short_clip_floor_spread(): + # extract_t (20) > total (7): evenly spaced timestamps, deduplicated, + # then pair-padded to an even count. + assert glm_sample_frame_indices(7, 30.0, 10.0) == [0, 1, 2, 3, 4, 5, 6, 6] + + +@pytest.mark.parametrize( + ("kwargs", "expected_len"), + [ + ({"target_fps": 0.5}, 16), + ({"max_frame_count": 16}, 16), + ({"target_fps": 8}, 240), # 30s * 8 = 240, under the 2048 cap + ], +) +def test_glm_sample_frame_indices_request_overrides(kwargs, expected_len): + indices = glm_sample_frame_indices(900, 30.0, 30.0, **kwargs) + assert len(indices) == expected_len + assert len(indices) % 2 == 0 + + +def test_sample_frames_fps_interval(video_processor): + indices = video_processor.sample_frames( + VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0) + ) + assert len(indices) == 60 + assert indices[0] == 0 + assert indices[-1] == 899 # uniform spread reaches the final frame + + # Request overrides reach the sampler through both kwarg spellings. The + # 15-frame spread is odd, so pair-padding duplicates the last frame. + assert ( + len( + video_processor.sample_frames( + VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0), + fps=0.5, + ) + ) + == 16 + ) + assert ( + len( + video_processor.sample_frames( + VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0), + target_fps=0.5, + ) + ) + == 16 + ) + assert ( + len( + video_processor.sample_frames( + VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0), + max_frames=16, + ) + ) + == 16 + ) + + +def test_defaults_mirror_checkpoint_config(): + """Bare instantiation matches the checkpoint's ``processor_config.json`` + token-budget style defaults; ``size`` carries no budget.""" + image_processor = Glm5NextImageProcessor() + assert image_processor.patch_expand_factor == 1 + assert image_processor.min_image_tokens == 16 + assert image_processor.max_image_tokens == 8000 + assert image_processor.resize_mode == "pad" + + video_processor = Glm5NextVideoProcessor() + assert video_processor.patch_expand_factor == 1 + assert video_processor.min_image_tokens == 16 + assert video_processor.max_image_tokens == 240000 + assert video_processor.fps_interval == 2.0 + assert video_processor.max_frame_count_dynamic == 2048 + + +def test_token_budgets_drive_geometry(): + """The token bounds fully determine the pixel budget and the grid.""" + token_proc = Glm5NextImageProcessor(min_image_tokens=64, max_image_tokens=512) + img = Image.new("RGB", (400, 300)) + grid = token_proc(img, return_tensors="pt")["image_grid_thw"][0].tolist() + # Budgets 100,352..802,816 px: the 300x400 canvas ceils to (308, 420) + # without hitting either bound -> 22x30 patches. + assert grid == [1, 22, 30] + assert token_proc.get_number_of_image_patches(300, 400) == 22 * 30 + + +def test_missing_token_budgets_rejected(): + proc = Glm5NextImageProcessor(min_image_tokens=None) + with pytest.raises(ValueError, match="min_image_tokens"): + proc(Image.new("RGB", (400, 300)), return_tensors="pt") + + +def test_video_config_fields_land(): + """fps_interval / max_frame_count_dynamic from the dedicated config + shape sampling without any request overrides.""" + proc = Glm5NextVideoProcessor( + fps_interval=4, + max_frame_count_dynamic=32, + ) + indices = proc.sample_frames( + VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0) + ) + # 30s * 4 = 120 candidates, capped at 32 -> uniform spread of 32. + assert len(indices) == 32 + assert indices[0] == 0 + assert indices[-1] == 899 + + # Request overrides still win over the config values. + assert ( + len( + proc.sample_frames( + VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0), + fps=0.5, + ) + ) + == 16 + ) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 76626418e1..c7207e55aa 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -49,6 +49,29 @@ DEVICE = torch.device("cpu") +def test_sm70_qwen_gdn_warmup_uses_bound_forward_context(monkeypatch): + calls: list[str] = [] + + class BoundGDNLayer: + def _warmup_sm70_causal_conv1d_real_state(self) -> bool: + calls.append("bound") + return True + + monkeypatch.setattr( + qwen_gdn.envs, "VLLM_SM70_AUX_KERNEL_WARMUP", True, raising=False + ) + monkeypatch.setattr( + qwen_gdn.current_platform, + "is_device_capability", + lambda capability: capability == 70, + ) + + assert qwen_gdn._warmup_sm70_qwen_gdn_causal_conv1d( + {"unrelated": object(), "model.layers.0.linear_attn": BoundGDNLayer()} + ) + assert calls == ["bound"] + + def test_disabled_gdn_core_dump_avoids_cuda_runtime_query(monkeypatch): monkeypatch.delenv("VLLM_SM70_DUMP_GDN_CORE_DIR", raising=False) monkeypatch.delenv("VLLM_SM70_DUMP_GDN_GRAPH_BUFFERS", raising=False) diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py new file mode 100644 index 0000000000..baf25a0431 --- /dev/null +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -0,0 +1,331 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU tests for the kpool tail slot mapping (no GPU required). + +The kpool tail cache is a 1-block-per-request circular ring addressed by +``pos % kpool`` (``KpoolTailSpec`` / ``KpoolTailManager``: exactly one block +allocated per request, never grown, so only column 0 of its block table is +ever written; the rest stays zero-initialized). + +The generic per-group slot kernel cannot express that layout: it maps +``pos -> bt[req][pos // bs] * bs + pos % bs`` (``_compute_slot_mappings_kernel`` +in vllm/v1/worker/gpu/block_table.py), so every token at ``pos >= kpool`` +reads a zero column and collapses onto physical tail block 0. All concurrent +requests then share one ``kpool``-slot ring and corrupt each other's pool +compression. + +These tests pin that defect's arithmetic, verify the circular replacement +(``compute_kpool_tail_slot_mapping``), and mirror the tail kernels' index +math to show cross-request pollution before the fix and isolation after. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + KpoolTailMetadataBuilder, + compute_kpool_tail_slot_mapping, +) + +KPOOL = 4 + + +def make_tail_block_table(own_blocks, width=64): + """Tail-group block table as BlockTables produces it: column 0 holds the + request's single KpoolTailManager block, the remaining columns are never + written and stay zero.""" + bt = torch.zeros(len(own_blocks), width, dtype=torch.int32) + bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32) + return bt + + +def legacy_generic_tail_slots(block_table, query_start_loc, positions): + """Reference of the generic ``_compute_slot_mappings_kernel`` arithmetic + (block_table.py:305-313) applied to the tail group's table.""" + slots = [] + for req in range(block_table.shape[0]): + for i in range(query_start_loc[req], query_start_loc[req + 1]): + pos = int(positions[i]) + block_number = int(block_table[req, pos // KPOOL]) + slots.append(block_number * KPOOL + pos % KPOOL) + return torch.tensor(slots, dtype=torch.int64) + + +def circular_tail_slots( + slot_mapping, block_table, query_start_loc, positions, num_actual, num_reqs +): + return compute_kpool_tail_slot_mapping( + slot_mapping, + block_table, + query_start_loc, + positions, + num_actual, + num_reqs, + KPOOL, + ) + + +def make_batch(per_req_positions, padded_len=None): + positions = torch.cat( + [torch.tensor(p, dtype=torch.int64) for p in per_req_positions] + ) + num_actual = positions.numel() + num_reqs = len(per_req_positions) + lens = [len(p) for p in per_req_positions] + qsl = torch.zeros(num_reqs + 1, dtype=torch.int64) + torch.cumsum(torch.tensor(lens, dtype=torch.int64), 0, out=qsl[1:]) + if padded_len is None: + padded_len = num_actual + slot_mapping = torch.full((padded_len,), -1, dtype=torch.int64) + return positions, qsl, slot_mapping, num_actual, num_reqs + + +def test_legacy_generic_mapping_collapses_onto_block_zero(): + """The bug: with the manager's 1-column block table, the generic kernel + maps every pos >= kpool onto tail block 0, and distinct requests collide.""" + own_blocks = [5, 9] + per_req = [list(range(10)), list(range(12))] # prompts of len 10 and 12 + positions, qsl, _, num_actual, num_reqs = make_batch(per_req) + bt = make_tail_block_table(own_blocks) + + legacy = legacy_generic_tail_slots(bt, qsl, positions) + + # Every token at pos >= kpool resolves to block 0, not the request's own. + off = 0 + for req, prompt in enumerate(per_req): + req_slots = legacy[off : off + len(prompt)] + for pos in range(len(prompt)): + slot = int(req_slots[pos]) + if pos >= KPOOL: + assert slot // KPOOL == 0, ( + f"expected collapse onto block 0 at pos {pos}" + ) + assert slot // KPOOL != own_blocks[req] + off += len(prompt) + + # The two requests share ring slots -> cross-request pollution. + a_slots = set(legacy[: len(per_req[0])].tolist()) + b_slots = set(legacy[len(per_req[0]) :].tolist()) + assert a_slots & b_slots, "legacy mapping must collide across requests" + + +def test_circular_mapping_isolates_requests(): + """The fix: every token lands in its own request's block at pos % kpool, + and no slot is ever shared by two different requests (slots do recur + within a request every kpool positions -- that is the circular design).""" + own_blocks = [5, 9] + per_req = [list(range(10)), list(range(12))] + positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) + bt = make_tail_block_table(own_blocks) + + out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) + + off = 0 + per_req_slots = [] + for req, prompt in enumerate(per_req): + req_slots = set() + for pos in range(len(prompt)): + slot = int(out[off + pos]) + assert slot // KPOOL == own_blocks[req], ( + f"req {req} pos {pos} left its tail block" + ) + assert slot % KPOOL == pos % KPOOL + req_slots.add(slot) + per_req_slots.append(req_slots) + off += len(prompt) + assert not per_req_slots[0] & per_req_slots[1] + + +@pytest.mark.parametrize("prompt_len", [1, 2, 3, 4]) +def test_circular_mapping_matches_generic_for_short_requests(prompt_len): + """For pos < kpool the generic kernel already picks the own block, so the + two mappings agree while every position fits the request's first block + (single-request behavior is unchanged).""" + own_blocks = [7] + per_req = [list(range(prompt_len))] + positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) + bt = make_tail_block_table(own_blocks) + + legacy = legacy_generic_tail_slots(bt, qsl, positions) + out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) + assert torch.equal(out, legacy) + + +def test_circular_mapping_preserves_padding_and_empty_batch(): + own_blocks = [5, 9] + per_req = [list(range(10)), list(range(12))] + padded_len = sum(len(p) for p in per_req) + 8 + positions, qsl, slot_mapping, num_actual, num_reqs = make_batch( + per_req, padded_len=padded_len + ) + bt = make_tail_block_table(own_blocks) + + out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) + assert out.shape == slot_mapping.shape + assert torch.equal(out[num_actual:], torch.full_like(out[num_actual:], -1)) + + empty = circular_tail_slots(slot_mapping, bt, qsl, positions[:0], 0, num_reqs) + assert torch.equal(empty, slot_mapping) + + +def make_common_metadata(per_req_positions, own_blocks, with_positions=True): + positions, qsl, slot_mapping, num_actual, num_reqs = make_batch( + per_req_positions, padded_len=sum(len(p) for p in per_req_positions) + 4 + ) + bt = make_tail_block_table(own_blocks) + seq_lens = torch.tensor( + [max(p) + 1 if p else 1 for p in per_req_positions], dtype=torch.int64 + ) + return CommonAttentionMetadata( + query_start_loc=qsl, + query_start_loc_cpu=qsl.clone(), + seq_lens=seq_lens, + num_reqs=num_reqs, + num_actual_tokens=num_actual, + max_query_len=max((len(p) for p in per_req_positions), default=1), + max_seq_len=int(seq_lens.max()) if num_reqs else 1, + block_table_tensor=bt, + slot_mapping=slot_mapping, + positions=positions if with_positions else None, + ) + + +def make_tail_builder(block_size=KPOOL): + builder = object.__new__(KpoolTailMetadataBuilder) + builder.kv_cache_spec = SimpleNamespace(block_size=block_size) + return builder + + +def test_builder_build_uses_circular_mapping(): + per_req = [list(range(10)), list(range(12))] + own_blocks = [5, 9] + cam = make_common_metadata(per_req, own_blocks) + meta = KpoolTailMetadataBuilder.build(make_tail_builder(), 0, cam) + + out = meta.slot_mapping + off = 0 + for req, prompt in enumerate(per_req): + for pos in range(len(prompt)): + slot = int(out[off + pos]) + assert slot // KPOOL == own_blocks[req] + assert slot % KPOOL == pos % KPOOL + off += len(prompt) + # Padding tail of the buffer keeps the -1 sentinel. + assert torch.equal( + out[cam.num_actual_tokens :], torch.full_like(out[cam.num_actual_tokens :], -1) + ) + + +def test_builder_build_falls_back_without_positions(): + """Capture / dummy builds without positions keep the generic mapping.""" + per_req = [list(range(10))] + cam = make_common_metadata(per_req, [5], with_positions=False) + meta = KpoolTailMetadataBuilder.build(make_tail_builder(), 0, cam) + assert meta.slot_mapping is cam.slot_mapping + + +# --------------------------------------------------------------------------- +# Index-level mirror of the tail kernels: seed / stash / pool completion +# (addressing replicated from kpool_compress.py's Triton kernels). +# --------------------------------------------------------------------------- + + +class TailRingMirror: + """Mirror of _kpool_tail_seed_kernel / _kpool_decode_update_batched_kernel + addressing: block = tail_slot // kpool, ring offset = pos % kpool; a pool + completing at pos reads ring slots (pool_start + s) % kpool and uses the + current token's own K/score for the last member.""" + + def __init__(self, num_blocks, kpool=KPOOL): + self.kpool = kpool + self.k = torch.full((num_blocks, kpool, 3), float("nan")) + self.s = torch.full((num_blocks, kpool, 3), float("nan")) + + def stash(self, tail_slot, pos, k, s): + blk, off = tail_slot // self.kpool, pos % self.kpool + self.k[blk, off] = k + self.s[blk, off] = s + + seed = stash # the seed kernel writes with the same addressing + + def complete(self, tail_slot, pos, k, s): + blk = tail_slot // self.kpool + start = pos - (self.kpool - 1) + kk = torch.stack( + [self.k[blk, (start + i) % self.kpool] for i in range(self.kpool)] + ) + ss = torch.stack( + [self.s[blk, (start + i) % self.kpool] for i in range(self.kpool)] + ) + kk[-1], ss[-1] = k, s # is_current for the completing token + w = torch.softmax(ss, dim=0) + return (kk * w).sum(0) + + +def token_kv(req, pos): + k = torch.tensor([pos + 100.0 * req, pos + 0.5, 2.0 * pos + 0.25]) + s = torch.tensor([0.1 * (pos + 1) + req, 0.2 * pos, 0.05 * pos]) + return k, s + + +def tail_slot_for(mapping, req, pos, own_block): + if mapping == "legacy": + bt_val = own_block if pos < KPOOL else 0 + return bt_val * KPOOL + pos % KPOOL + return own_block * KPOOL + pos % KPOOL + + +def run_scenario(mapping, interleave): + """Requests A (block 5, prompt len 9) and B (block 9, prompt len 11) + decode concurrently; returns A's boundary pool [8, 9, 10, 11].""" + ring = TailRingMirror(num_blocks=16) + blocks = {"A": 5, "B": 9} + prompts = {"A": 9, "B": 11} + + def slot(req, pos): + return tail_slot_for(mapping, 0 if req == "A" else 1, pos, blocks[req]) + + # Prefill: seed each request's trailing incomplete pool. + for req, L in prompts.items(): + for pos in range(L - (L % KPOOL or KPOOL), L): + if pos < 0: + continue + ring.stash(slot(req, pos), pos, *token_kv(0 if req == "A" else 1, pos)) + + # Decode: A emits pos 9, 10, 11; B emits 11, 12, 13. `interleave` + # processes B before A within a step, which is exactly what concurrent + # Triton programs do when they share tail block 0. + order = ["B", "A"] if interleave else ["A", "B"] + decode = {"A": [9, 10, 11], "B": [11, 12, 13]} + result = None + for step in range(3): + for req in order: + pos = decode[req][step] + k, s = token_kv(0 if req == "A" else 1, pos) + if pos % KPOOL == KPOOL - 1: + pool = ring.complete(slot(req, pos), pos, k, s) + if req == "A": + result = pool + ring.stash(slot(req, pos), pos, k, s) + return result + + +def test_interleaved_decode_pollution_legacy_vs_circular(): + """Ground truth: request A decoding alone (its ring used exclusively).""" + torch.manual_seed(0) + ground_truth = run_scenario("circular", interleave=False) + + legacy = run_scenario("legacy", interleave=True) + circular = run_scenario("circular", interleave=True) + + # The old mapping lets request B's tokens into A's ring: A's boundary + # pool is compressed from 2 of B's tokens -> wrong. + assert not torch.allclose(legacy, ground_truth), ( + f"legacy mapping unexpectedly clean: {legacy} vs {ground_truth}" + ) + + # The circular mapping keeps the rings isolated under interleaving. + torch.testing.assert_close(circular, ground_truth) diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index fb2e195796..8627c64663 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -738,6 +738,44 @@ def test_sm70_splitd_d256_loader_requires_exact_ops(monkeypatch): assert flash_v100._get_sm70_splitd_d256_ops() == (dense, paged, splitkv3) +def test_sm70_splitd_d256_loader_accepts_explicit_sidecar(monkeypatch): + import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 + + fake_interface = types.ModuleType("vllm.vllm_flash_attn.flash_attn_interface") + fake_package = types.ModuleType("vllm.vllm_flash_attn") + fake_package.__dict__["flash_attn_interface"] = fake_interface + monkeypatch.setitem(sys.modules, "vllm.vllm_flash_attn", fake_package) + monkeypatch.setitem( + sys.modules, + "vllm.vllm_flash_attn.flash_attn_interface", + fake_interface, + ) + + namespace = SimpleNamespace() + loaded: list[str] = [] + + def load_library(path: str) -> None: + loaded.append(path) + namespace.sm70_d256_splitd_n32_dense_fwd = "dense" + namespace.sm70_d256_splitd_n32_paged_fwd = "paged" + + fake_ops = SimpleNamespace( + _vllm_fa2_C=namespace, + load_library=load_library, + ) + monkeypatch.setattr(flash_v100, "torch", SimpleNamespace(ops=fake_ops)) + monkeypatch.setenv("VLLM_SM70_FA2_D256_LIBRARY", "/tmp/stable-fa2.so") + monkeypatch.setattr(flash_v100, "_sm70_splitd_d256_ops_checked", False) + monkeypatch.setattr(flash_v100, "_sm70_splitd_d256_ops", None) + + assert flash_v100._get_sm70_splitd_d256_ops() == ( + "dense", + "paged", + None, + ) + assert loaded == ["/tmp/stable-fa2.so"] + + def test_sm70_d256_gqa_architecture_loader_is_optional(monkeypatch): import vllm.v1.attention.backends.flash_attn_v100 as flash_v100 @@ -891,7 +929,7 @@ def test_prefill_d256_gqa_architecture_policy_is_shape_family_bounded(monkeypatc monkeypatch.delenv(name, raising=False) envs.disable_envs_cache() - for kv_len in range(40000, 128001, 8000): + for kv_len in range(16000, 256001, 8000): family_key = torch.empty( (1, kv_len, 1, 256), dtype=torch.float16, device="meta" ) @@ -904,6 +942,16 @@ def test_prefill_d256_gqa_architecture_policy_is_shape_family_bounded(monkeypatc softmax_scale=0.0625, architecture_op=object(), ) + first_chunk_key = torch.empty((1, 8000, 1, 256), dtype=torch.float16, device="meta") + assert not flash_v100._should_use_prefill_d256_gqa_architecture( + query, + first_chunk_key, + torch.empty_like(first_chunk_key), + max_seqlen_q=8000, + max_seqlen_k=8000, + softmax_scale=0.0625, + architecture_op=object(), + ) assert not flash_v100._should_use_prefill_d256_gqa_architecture( query[:, :7999], key, @@ -924,10 +972,10 @@ def test_prefill_d256_gqa_architecture_policy_is_shape_family_bounded(monkeypatc ) assert not flash_v100._should_use_prefill_d256_gqa_architecture( query, - key[:, :44000], - value[:, :44000], + key[:, :12000], + value[:, :12000], max_seqlen_q=8000, - max_seqlen_k=44000, + max_seqlen_k=12000, softmax_scale=0.0625, architecture_op=object(), ) @@ -1655,7 +1703,11 @@ def fake_decode( assert torch.all(output == 1) -def test_flash_v100_dflash2_grouped_verify_uses_original_request_metadata(): +@pytest.mark.parametrize("query_len", [8, 16]) +@pytest.mark.parametrize("page_size", [3296, 3456]) +def test_flash_v100_dflash2_grouped_verify_uses_original_request_metadata( + query_len: int, page_size: int +): from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl impl = FlashAttnV100Impl( @@ -1690,21 +1742,21 @@ def grouped_verify( original_block_table = torch.tensor([[7, 3]], dtype=torch.int32) original_seq_lens = torch.tensor([2056], dtype=torch.int32) attn_metadata = SimpleNamespace( - num_actual_tokens=8, + num_actual_tokens=query_len, causal=True, is_dflash_selector_target=True, max_model_len=32768, - query_start_loc=torch.tensor([0, 8], dtype=torch.int32), + query_start_loc=torch.tensor([0, query_len], dtype=torch.int32), seq_lens=original_seq_lens, block_table=original_block_table, - smallq_decode_block_table=torch.zeros((8, 2), dtype=torch.int32), - smallq_decode_seq_lens=torch.arange(2049, 2057, dtype=torch.int32), - smallq_query_start_loc=torch.tensor([0, 8], dtype=torch.int32), + smallq_decode_block_table=torch.zeros((query_len, 2), dtype=torch.int32), + smallq_decode_seq_lens=torch.arange(2057 - query_len, 2057, dtype=torch.int32), + smallq_query_start_loc=torch.tensor([0, query_len], dtype=torch.int32), ) layer = SimpleNamespace(_k_scale_float=0.5, _v_scale_float=2.0) - query = torch.zeros((8, 6, 256), dtype=torch.float16) + query = torch.zeros((query_len, 6, 256), dtype=torch.float16) output = torch.zeros_like(query) - key_cache = torch.zeros((2, 3296, 1, 256), dtype=torch.uint8) + key_cache = torch.zeros((2, page_size, 1, 256), dtype=torch.uint8) value_cache = torch.zeros_like(key_cache) attn_metadata.max_model_len = 8192 @@ -1713,7 +1765,7 @@ def grouped_verify( key_cache, value_cache, attn_metadata, - num_query_tokens=8, + num_query_tokens=query_len, ) attn_metadata.max_model_len = 32768 result = impl._flash_v100_small_query_prefill_as_decode( diff --git a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py new file mode 100644 index 0000000000..31099ed1e2 --- /dev/null +++ b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU tests for the decode top-k row->token mapping (no GPU required). + +On non-uniform decode batches (``requires_padding`` -- mixed plain-decode and +spec-verify requests, or variable MTP verify lens; taken on Hopper where the +varlen/flatten logits path is unavailable), the pool-topk rows follow the +PADDED ``[batch_size, next_n]`` grid: row ``(b, t)`` is flat decode token +``offset_b + t``. The former inline ``dec_seq = positions[:n] + 1`` +(``n = batch_size * next_n``) indexes the flat per-token layout with padded +coordinates, so rows after the first non-uniform request read another +request's positions and rows past the decode region read prefill tokens -- +``expand_pools_and_append_tail`` then anchors the tail at a foreign length. + +These tests pin the defect with the exact production arithmetic and verify +the layout-aware replacement (``_decode_topk_seq_lens``), including the tail +expansion consequences via the pure-torch expand/append pair that the fused +kernel is documented to replicate. +""" + +import torch + +# Bootstrap the glm5next package before entering the indexer module: its +# kpool_compress import runs glm5next/__init__, which pulls model -> +# attention -> back into sparse_attn_indexer_kpool (attention.py imports the +# class at module scope). Production always enters via attention.py first. +# isort: off +import vllm.models.glm5next # noqa: F401 +from vllm.models.glm5next.nvidia.ops.kpool_compress import ( # noqa: E402 + append_tail_to_topk, + expand_pools_to_tokens, +) +# isort: on + +from vllm.model_executor.layers.sparse_attn_indexer_kpool import ( + _decode_topk_seq_lens, +) + +KPOOL = 4 +TOPK_TOKENS = 16 +SELECT_K = TOPK_TOKENS // KPOOL + + +def make_non_uniform_batch(): + """3 requests: plain decode (1 token), MTP verify (4), adaptive verify (3). + + Flat decode positions (production layout: decode tokens first, then any + prefill tokens of the same batch): + req0: [30] (context len 30) + req1: [100..103] (verify at context len 100) + req2: [7, 8, 9] (verify at context len 7) + Followed by 4 prefill tokens at positions 555..558 so the flat tensor is + at least ``n = 3 * 4`` long, as it is in a real mixed batch. + """ + decode_lens = torch.tensor([1, 4, 3], dtype=torch.int64) + per_req_positions = [[30], [100, 101, 102, 103], [7, 8, 9]] + flat_decode = [p for req in per_req_positions for p in req] + positions = torch.tensor(flat_decode + [555, 556, 557, 558], dtype=torch.int64) + return decode_lens, per_req_positions, positions + + +def expected_row_seq_lens(per_req_positions, batch_size, next_n): + """Ground truth: row (b, t) -> pos + 1 for real rows, 0 for pad rows.""" + out = torch.zeros(batch_size * next_n, dtype=torch.int32) + for b, req in enumerate(per_req_positions): + for t, pos in enumerate(req): + out[b * next_n + t] = pos + 1 + return out + + +def test_legacy_flat_layout_misaligns_and_bleeds(): + """The bug: ``positions[:n] + 1`` with padded-row coordinates reads other + requests' positions and prefill positions.""" + decode_lens, per_req, positions = make_non_uniform_batch() + batch_size = decode_lens.shape[0] + next_n = int(decode_lens.max()) + n = batch_size * next_n + assert n == 12 + + legacy = positions[:n].to(torch.int32) + 1 + expected = expected_row_seq_lens(per_req, batch_size, next_n) + + # Row (1, 0): req1's first verify token (pos 100, seq 101) reads flat + # index 4 -> req1's LAST verify token (pos 103). Its true tail token is + # dropped and the tail anchors 3 tokens late. + assert int(legacy[4]) == 104 and int(expected[4]) == 101 + + # Rows (2, *): req2 starts at flat offset 5, but padded coordinates point + # at flat indices 8..11 -- the batch's PREFILL positions. + for t in range(3): + assert int(legacy[2 * next_n + t]) >= 556, ( + "legacy row (2, t) should read prefill positions" + ) + assert not torch.equal(legacy, expected) + + +def test_helper_uniform_layout_matches_flat_slice(): + """Uniform batches keep the flat shortcut (zero behavior/perf change).""" + per_req = [[200, 201, 202, 203], [50, 51, 52, 53]] + positions = torch.tensor([p for r in per_req for p in r], dtype=torch.int64) + decode_lens = torch.tensor([4, 4], dtype=torch.int64) + out = _decode_topk_seq_lens(positions, decode_lens, 8, 2, 4, requires_padding=False) + assert torch.equal(out, positions[:8].to(torch.int32) + 1) + + +def test_helper_padded_layout_per_row(): + """Non-uniform batches map every padded row to its own token's position; + pad rows get 0 (empty tail).""" + decode_lens, per_req, positions = make_non_uniform_batch() + out = _decode_topk_seq_lens(positions, decode_lens, 8, 3, 4, requires_padding=True) + expected = expected_row_seq_lens(per_req, 3, 4) + assert torch.equal(out, expected) + # Pad rows (0, 1..3) and (2, 3) collapse to 0 -> no tail appended. + assert out[1] == 0 and out[2] == 0 and out[3] == 0 and out[11] == 0 + + +def expand_tail_region(dec_seq): + """Run the production pure-torch expand + append pair (the fused kernel + replicates it exactly on the identity path) and return the tail columns + [TOPK_TOKENS, TOPK_TOKENS + KPOOL - 1).""" + rows = dec_seq.shape[0] + pool_ids = torch.arange(SELECT_K, dtype=torch.int64).expand(rows, SELECT_K) + valid = torch.ones_like(pool_ids, dtype=torch.bool) + expanded = expand_pools_to_tokens(pool_ids, valid, TOPK_TOKENS, KPOOL) + seq_lens = dec_seq.to(torch.int32) + pool_lens = (seq_lens // KPOOL).to(torch.int32) + out = append_tail_to_topk(expanded, seq_lens, pool_lens, KPOOL) + return out[:, TOPK_TOKENS:] + + +def test_tail_expansion_legacy_vs_fixed(): + """End-to-end tail consequence: the fixed mapping appends exactly the + request's trailing incomplete pool; the legacy flat slice drops real tail + tokens and emits indices far past the request's own sequence.""" + decode_lens, per_req, positions = make_non_uniform_batch() + n = 12 + fixed = _decode_topk_seq_lens( + positions, decode_lens, 8, 3, 4, requires_padding=True + ) + legacy = positions[:n].to(torch.int32) + 1 + + fixed_tail = expand_tail_region(fixed) + legacy_tail = expand_tail_region(legacy) + + # Fixed: row (1, 0) (seq 101) keeps its single tail token 100; row (2, 2) + # (seq 10) keeps its full trailing pool [8, 9]. + assert fixed_tail[4, 0].item() == 100 + assert fixed_tail[10, :2].tolist() == [8, 9] + + # Legacy: row (1, 0) loses its tail token (anchored at 104, count 0)... + assert legacy_tail[4, 0].item() == -1 + # ...and row (2, 2) reads prefill position 557 -> tail indices 556/557, + # way past req2's 10-token sequence -> out-of-bounds block-table reads. + assert legacy_tail[10, :2].tolist() == [556, 557] + + assert not torch.equal(legacy_tail, fixed_tail) diff --git a/tests/v1/core/test_glm5_kv_cache_layout.py b/tests/v1/core/test_glm5_kv_cache_layout.py new file mode 100644 index 0000000000..5f79728d69 --- /dev/null +++ b/tests/v1/core/test_glm5_kv_cache_layout.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm.v1.core import kv_cache_utils +from vllm.v1.kv_cache_interface import ( + KpoolTailSpec, + MambaSpec, + MLAAttentionSpec, +) + + +def test_glm53_pp2_kpool_tail_shares_indexer_storage(monkeypatch): + monkeypatch.delenv("VLLM_PP_LAYER_PARTITION", raising=False) + model_config = SimpleNamespace( + max_model_len=32768, + get_total_num_hidden_layers=lambda: 45, + ) + vllm_config = SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace(pipeline_parallel_size=2), + cache_config=SimpleNamespace( + mamba_cache_mode="align", + num_gpu_blocks_override=None, + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=2048, + disable_hybrid_kv_cache_manager=False, + ), + ) + + block_size = 1152 + mamba_block_size = 4096 + mamba_spec = MambaSpec( + shapes=((3, 2048), (3, 2048), (3, 2048), (16, 128, 128)), + dtypes=(torch.float16, torch.float16, torch.float16, torch.float32), + block_size=mamba_block_size, + ) + specs = {} + for layer_idx in range(45): + prefix = f"model.layers.{layer_idx}.self_attn" + if layer_idx % 4 == 3: + specs[f"{prefix}.mla_attn"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=512, + dtype=torch.float16, + ) + specs[f"{prefix}.indexer.k_cache"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + compress_ratio=4, + ) + specs[f"{prefix}.indexer.tail_cache"] = KpoolTailSpec( + block_size=4, + num_kv_heads=1, + head_size=256, + head_size_v=0, + dtype=torch.bfloat16, + sliding_window=4, + ) + else: + specs[prefix] = mamba_spec + + groups = kv_cache_utils.get_kv_cache_groups(vllm_config, specs) + layout = kv_cache_utils._glm5_next_tensor_layout(groups) + assert layout is not None + _, mamba_groups, mla_names, idx_names, mla_page, idx_page, tail_names, _ = layout + assert len(mamba_groups) == 4 + assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8] + assert all( + group.kv_cache_spec.block_size == mamba_block_size for group in mamba_groups + ) + assert len(mla_names) == len(idx_names) == len(tail_names) == 11 + participating_block_sizes = [ + group.kv_cache_spec.block_size + for group in groups + if group.kv_cache_spec.prefix_cacheable + ] + assert min(participating_block_sizes) == block_size + assert min(group.kv_cache_spec.block_size for group in groups) == 4 + + bytes_per_block = 11 * (mla_page + idx_page) + cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + vllm_config, groups, bytes_per_block * 100 + ) + assert cache_config.num_blocks == 100 + assert len(cache_config.kv_cache_tensors) == 22 + for idx_name, tail_name in zip(idx_names, tail_names): + tensor = next( + t for t in cache_config.kv_cache_tensors if idx_name in t.shared_by + ) + assert tensor.shared_by == [idx_name, tail_name] diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 1ec844fead..52542457f5 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -183,6 +183,36 @@ def new_mamba_spec( ) +def test_small_draft_sliding_window_bucket_absorbs_hybrid_padding(): + kv_cache_specs: dict[str, KVCacheSpec] = {} + full_spec = new_kv_cache_spec() + mamba_spec = new_mamba_spec(page_size_padded=full_spec.page_size_bytes) + sliding_spec = new_sliding_window_spec(sliding_window=2048) + kv_cache_specs.update({f"full.{i}": full_spec for i in range(16)}) + kv_cache_specs.update({f"mamba.{i}": mamba_spec for i in range(48)}) + kv_cache_specs.update({f"draft.{i}": sliding_spec for i in range(5)}) + + groups = kv_cache_utils._get_kv_cache_groups_uniform_page_size(kv_cache_specs) + + assert len(groups) == 9 + assert [len(group.layer_names) for group in groups] == [8] * 8 + [5] + + +def test_sliding_window_padding_preserves_generic_singleton_grouping(): + kv_cache_specs: dict[str, KVCacheSpec] = {} + full_spec = new_kv_cache_spec() + mamba_spec = new_mamba_spec(page_size_padded=full_spec.page_size_bytes) + sliding_spec = new_sliding_window_spec(sliding_window=2048) + kv_cache_specs.update({f"full.{i}": full_spec for i in range(15)}) + kv_cache_specs.update({f"mamba.{i}": mamba_spec for i in range(45)}) + kv_cache_specs.update({"draft.0": sliding_spec}) + + groups = kv_cache_utils._get_kv_cache_groups_uniform_page_size(kv_cache_specs) + + assert len(groups) == 13 + assert [len(group.layer_names) for group in groups] == [5] * 12 + [1] + + def test_resolve_kv_cache_block_sizes_for_aligned_mamba_group(): """Aligned Mamba keeps fine-grained hashes after page-size unification.""" kv_cache_config = KVCacheConfig( @@ -214,6 +244,26 @@ def test_resolve_kv_cache_block_sizes_for_aligned_mamba_group(): ) == (1648, 16) +@pytest.mark.parametrize( + ("layer_counts", "expected_group_size"), + [ + ([1, 12, 12, 12, 36], 12), + ([1, 20, 30], 10), + ([1, 64, 64], 16), + ([1, 5, 7], 1), + ([3, 7], 3), + ([5, 6], 6), + ], +) +def test_select_hybrid_kv_cache_group_size( + layer_counts: list[int], expected_group_size: int +): + assert ( + kv_cache_utils._select_hybrid_kv_cache_group_size(layer_counts) + == expected_group_size + ) + + @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_none_hash(monkeypatch, hash_fn): import vllm.v1.core.kv_cache_utils diff --git a/tests/v1/core/test_qwen4_exp_kv_cache.py b/tests/v1/core/test_qwen4_exp_kv_cache.py new file mode 100644 index 0000000000..4b777dc2e5 --- /dev/null +++ b/tests/v1/core/test_qwen4_exp_kv_cache.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm.models.qwen4_exp.common.qsa_cache import QSAStateBackend +from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator +from vllm.v1.core.kv_cache_utils import ( + _get_csa_linear_tensor_layout, + generate_scheduler_kv_cache_config, + get_kv_cache_config_from_groups, + get_kv_cache_groups, +) +from vllm.v1.core.single_type_kv_cache_manager import CircularBufferManager +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + FullAttentionSpec, + MambaSpec, + MLAAttentionSpec, +) +from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache +from vllm.v1.worker.utils import AttentionGroup + + +class _ModelConfig: + max_model_len = 8192 + + def get_num_kv_heads(self, parallel_config) -> int: + del parallel_config + return 1 + + def get_total_num_hidden_layers(self) -> int: + return 8 + + +def _vllm_config(): + return SimpleNamespace( + model_config=_ModelConfig(), + parallel_config=SimpleNamespace(pipeline_parallel_size=1), + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + cache_config=SimpleNamespace( + num_gpu_blocks_override=None, + mamba_cache_mode="none", + ), + ) + + +def _qwen4_exp_cache_specs(): + specs = {} + for layer in (3, 7): + prefix = f"model.layers.{layer}.self_attn" + specs[prefix] = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=256, + head_size_v=256, + dtype=torch.float16, + ) + specs[f"{prefix}.compressed"] = MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.float16, + compress_ratio=4, + ) + specs[f"{prefix}.compressor_state"] = CircularBufferSpec( + block_size=4, + num_kv_heads=1, + head_size=128, + head_size_v=0, + dtype=torch.float16, + ) + + for layer in (0, 1, 2, 4, 5, 6): + specs[f"model.layers.{layer}.linear_attn"] = MambaSpec( + block_size=16, + shapes=((1, 64),), + dtypes=(torch.float16,), + ) + specs["model.layers.2.ple"] = MambaSpec( + block_size=16, + shapes=((1, 64),), + dtypes=(torch.float16,), + tp_replicated=True, + ) + return specs + + +def test_qwen4_exp_csa_linear_cache_layout() -> None: + groups = get_kv_cache_groups(_vllm_config(), _qwen4_exp_cache_specs()) + layout = _get_csa_linear_tensor_layout(groups) + + assert layout is not None + assert [len(group.layer_names) for group in groups] == [4, 2, 2, 2, 2, 1] + assert len(layout.main_kv_names) == 2 + assert len(layout.compressed_names) == 2 + assert len(layout.compressor_state_names) == 2 + assert len(layout.mamba_groups) == 4 + + cache_config = get_kv_cache_config_from_groups( + _vllm_config(), groups, available_memory=1 << 30 + ) + assert len(cache_config.kv_cache_tensors) == 4 + assert all(len(tensor.shared_by) >= 2 for tensor in cache_config.kv_cache_tensors) + + scheduler_config = generate_scheduler_kv_cache_config([cache_config]) + scheduler_config.num_blocks = 32 + assert isinstance( + scheduler_config.kv_cache_groups[0].kv_cache_spec, FullAttentionSpec + ) + assert isinstance( + scheduler_config.kv_cache_groups[1].kv_cache_spec, CircularBufferSpec + ) + coordinator = get_kv_cache_coordinator( + scheduler_config, + max_model_len=8192, + max_in_flight_tokens=128, + use_eagle=False, + enable_caching=False, + enable_kv_cache_events=False, + dcp_world_size=1, + pcp_world_size=1, + hash_block_size=4, + ) + assert isinstance(coordinator.single_type_managers[1], CircularBufferManager) + + prefix_coordinator = get_kv_cache_coordinator( + scheduler_config, + max_model_len=8192, + max_in_flight_tokens=128, + use_eagle=False, + enable_caching=True, + enable_kv_cache_events=False, + dcp_world_size=1, + pcp_world_size=1, + hash_block_size=4, + ) + assert all( + not isinstance(spec, CircularBufferSpec) + for spec, _, _ in prefix_coordinator.attention_groups + ) + + +def test_qwen4_exp_circular_cache_stores_keys_without_unused_values() -> None: + spec = CircularBufferSpec( + block_size=4, + num_kv_heads=1, + head_size=128, + head_size_v=0, + dtype=torch.float16, + ) + + assert spec.real_page_size_bytes == 4 * 128 * 2 + assert spec.max_memory_usage_bytes(_vllm_config()) == spec.page_size_bytes + + +def test_qwen4_exp_circular_manager_owns_one_block_per_request() -> None: + spec = CircularBufferSpec( + block_size=4, + num_kv_heads=1, + head_size=128, + head_size_v=0, + dtype=torch.float16, + ) + block_pool = BlockPool( + num_gpu_blocks=8, + enable_caching=False, + hash_block_size=spec.block_size, + ) + manager = CircularBufferManager( + spec, + block_pool=block_pool, + enable_caching=False, + kv_cache_group_id=0, + ) + + assert manager.get_num_blocks_to_allocate("req", 4096, (), 0, 4096) == 1 + blocks = manager.allocate_new_blocks("req", 4096, 4096) + assert len(blocks) == 1 + assert manager.req_to_blocks["req"] == blocks + assert manager.get_num_blocks_to_allocate("req", 8192, (), 4096, 8192) == 0 + assert manager.allocate_new_blocks("req", 8192, 8192) == [] + + +def test_qwen4_exp_compressed_qsa_reshape_uses_storage_block_size() -> None: + spec = MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.float16, + compress_ratio=4, + ) + num_blocks = 3 + raw = torch.empty(num_blocks * spec.page_size_bytes, dtype=torch.int8) + group = AttentionGroup( + QSAStateBackend, + ["compressed"], + spec, + kv_cache_group_id=0, + ) + + caches = _reshape_kv_cache( + attn_groups=[group], + kv_cache_raw_tensors={"compressed": raw}, + cache_dtype="auto", + kernel_block_sizes=[16], + shared_kv_cache_layers={}, + ) + + assert caches["compressed"].shape == (num_blocks, 4, 1, 128) + assert caches["compressed"].untyped_storage().data_ptr() == raw.data_ptr() diff --git a/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py b/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py index 23c58085e8..16fd480cc6 100644 --- a/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py +++ b/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py @@ -94,3 +94,53 @@ def test_split_managers_keep_exact_full_graph_shapes( assert draft_desc.cg_mode == CUDAGraphMode.FULL assert draft_desc.num_tokens == batch_size assert draft_desc.num_reqs == batch_size + + +def test_dflash_lookup_captures_b1_q8_and_q16(monkeypatch) -> None: + monkeypatch.setattr( + "vllm.v1.worker.gpu.cudagraph_utils.current_platform.get_global_graph_pool", + lambda: None, + ) + monkeypatch.setattr( + "vllm.v1.worker.gpu.cudagraph_utils.get_pp_group", + lambda: MagicMock(is_first_rank=True, is_last_rank=True), + ) + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_seqs=1), + parallel_config=SimpleNamespace( + tensor_parallel_size=4, + data_parallel_size=1, + ), + compilation_config=CompilationConfig( + cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY, + cudagraph_capture_sizes=[16], + max_cudagraph_capture_size=16, + ), + speculative_config=SimpleNamespace( + method="dflash", + ngram_assist=True, + num_speculative_tokens=15, + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace( + dflash_config={"block_size": 8, "selector_top_k": 16} + ) + ), + ), + ) + + manager = CudaGraphManager( + config, + torch.device("cpu"), + CUDAGraphMode.FULL_DECODE_ONLY, + decode_query_len=16, + ) + + assert manager.decode_query_lens == (8, 16) + assert manager._capture_sizes == [8, 16] + manager._graphs_captured = True + q8 = manager.dispatch(1, 8, 8) + q16 = manager.dispatch(1, 16, 16) + assert q8.cg_mode == CUDAGraphMode.FULL + assert q8.uniform_token_count == 8 + assert q16.cg_mode == CUDAGraphMode.FULL + assert q16.uniform_token_count == 16 diff --git a/tests/v1/spec_decode/test_ddtree_flex_attention.py b/tests/v1/spec_decode/test_ddtree_flex_attention.py index 4728babacb..e7a90e71e1 100644 --- a/tests/v1/spec_decode/test_ddtree_flex_attention.py +++ b/tests/v1/spec_decode/test_ddtree_flex_attention.py @@ -29,11 +29,11 @@ def test_ddtree_logical_mask_hides_siblings() -> None: ) assert mask.tolist() == [ - True, # history - True, # history - True, # root - True, # parent node 1 - True, # self node 2 + True, # history + True, # history + True, # root + True, # parent node 1 + True, # self node 2 False, # sibling node 3 False, # beyond tree ] diff --git a/tests/v1/spec_decode/test_ddtree_gpu_runner_positions.py b/tests/v1/spec_decode/test_ddtree_gpu_runner_positions.py index 5ddfa45fb4..2236e4c699 100644 --- a/tests/v1/spec_decode/test_ddtree_gpu_runner_positions.py +++ b/tests/v1/spec_decode/test_ddtree_gpu_runner_positions.py @@ -304,9 +304,7 @@ def copy_func( forward_context=forward_context, mamba_state_copy_funcs=(copy_func,), copy_bufs=copy_bufs, - ddtree_accepted_node_indices=torch.tensor( - [[0, 2, 5, 9]], dtype=torch.int32 - ), + ddtree_accepted_node_indices=torch.tensor([[0, 2, 5, 9]], dtype=torch.int32), ) assert calls == [(9, 1)] @@ -317,9 +315,7 @@ def test_update_states_after_model_execute_keeps_compact_state_selector() -> Non runner.speculative_config = object() runner.model_config = SimpleNamespace(is_hybrid=True) runner.cache_config = SimpleNamespace(mamba_cache_mode="none") - runner.num_accepted_tokens = SimpleNamespace( - gpu=torch.ones(1, dtype=torch.int32) - ) + runner.num_accepted_tokens = SimpleNamespace(gpu=torch.ones(1, dtype=torch.int32)) runner.spec_state_slot_selectors = SimpleNamespace( gpu=torch.ones(1, dtype=torch.int32) ) @@ -389,9 +385,7 @@ def test_compact_ddtree_drafter_context_moves_branch_to_prefix() -> None: scheduler_output = _scheduler_output(payload) scheduler_output.num_scheduled_tokens = {"r0": 6} scheduler_output.total_num_scheduled_tokens = 6 - scheduler_output.scheduled_spec_decode_tokens = { - "r0": list(payload.tree_token_ids) - } + scheduler_output.scheduled_spec_decode_tokens = {"r0": list(payload.tree_token_ids)} runner._compact_ddtree_drafter_context( hidden_states, diff --git a/tests/v1/spec_decode/test_ddtree_metadata.py b/tests/v1/spec_decode/test_ddtree_metadata.py index 3d58b3822c..6188ec103e 100644 --- a/tests/v1/spec_decode/test_ddtree_metadata.py +++ b/tests/v1/spec_decode/test_ddtree_metadata.py @@ -34,7 +34,8 @@ def test_single_request_metadata_indices() -> None: assert metadata.compact_logits_indices[0] == 6 assert metadata.compact_logits_indices[1] == 7 assert metadata.node_compact_indices == tuple( - range(1, tree.non_root_nodes[-1].index + 1)) + range(1, tree.non_root_nodes[-1].index + 1) + ) assert metadata.edge_parent_compact_indices[0] == 0 assert metadata.edge_parent_compact_indices[1] == 1 assert metadata.tree_position_ids[:3] == (7, 8, 9) diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 98cda4fe8c..6cf705973a 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -25,6 +25,7 @@ VocabParallelEmbedding, _sm70_dflash2_dense_order_topk, _sm70_dflash2_rerank_output_buffers, + _sm70_dflash2_use_dense_order, ) from vllm.model_executor.models.dflash_sm70 import ( DFLASH_SM70_GATE_UP_INPUT_SCALE, @@ -63,8 +64,8 @@ def test_dflash2_gdn_fastpaths_are_default_off(monkeypatch): names = ( - "VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY", "VLLM_SM70_DFLASH2_QPN8_RERANK", + "VLLM_SM70_DFLASH2_QPN8_ALLOW_CANDIDATE_ORDER", "VLLM_SM70_DFLASH2_VERIFY_FASTPATH", "VLLM_SM70_DFLASH2_FUSED_GDN_METADATA", "VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW", @@ -87,6 +88,19 @@ def test_dflash2_gdn_fastpaths_are_default_off(monkeypatch): envs.disable_envs_cache() +def test_dflash2_grouped_verify_is_default_on_with_rollback(monkeypatch): + name = "VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY" + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + try: + assert getattr(envs, name) + monkeypatch.setenv(name, "0") + envs.disable_envs_cache() + assert not getattr(envs, name) + finally: + envs.disable_envs_cache() + + def test_sm70_tp4_push_allreduce_is_default_on_with_rollback(monkeypatch): monkeypatch.delenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", raising=False) envs.disable_envs_cache() @@ -343,6 +357,7 @@ def init_base(self, _vllm_config, device): self.max_num_reqs = 2 self.num_query_per_req = 8 self.num_speculative_steps = 7 + self.draft_block = 7 self.vocab_size = 31 self.draft_tokens = torch.empty((2, 7), dtype=torch.int64, device=device) self.draft_logits = draft_logits @@ -927,6 +942,23 @@ def test_qpn8_rerank_restores_dense_vocab_tie_order(): assert torch.equal(actual_ids, expected_ids + vocab_start) +def test_qpn8_candidate_order_requires_explicit_experimental_opt_in(monkeypatch): + monkeypatch.setattr(envs, "VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER", False) + monkeypatch.setattr( + envs, + "VLLM_SM70_DFLASH2_QPN8_ALLOW_CANDIDATE_ORDER", + False, + ) + assert _sm70_dflash2_use_dense_order() + + monkeypatch.setattr( + envs, + "VLLM_SM70_DFLASH2_QPN8_ALLOW_CANDIDATE_ORDER", + True, + ) + assert not _sm70_dflash2_use_dense_order() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("selector_k", [16, 20]) @pytest.mark.parametrize("num_rows", [1, 7, 8]) @@ -1115,6 +1147,7 @@ def test_dflash_intermediate_prefill_materializes_context_without_query(monkeypa speculator.max_model_len = 256 speculator.num_query_per_req = 8 speculator.num_speculative_steps = 7 + speculator.draft_block = 7 speculator.max_num_reqs = 1 speculator.max_num_tokens = 8 speculator.hidden_states = torch.zeros(8, 3) diff --git a/tests/v1/spec_decode/test_dflash2_lookup.py b/tests/v1/spec_decode/test_dflash2_lookup.py new file mode 100644 index 0000000000..13001125d8 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash2_lookup.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.triton_utils import triton +from vllm.v1.worker.gpu.spec_decode.dflash2.lookup import ( + _point_mass_draft_logits_kernel, + fuse_draft, + suffix_lookup, +) +from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( + _prepare_lookup_controller_flags_kernel, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA is required for LABD kernels" +) + + +def _reference_lookup( + tokens: list[int], k: int, nmin: int, nmax: int +) -> tuple[list[int], int, int]: + suffix_end = len(tokens) - 1 + best = (-1, -1) + for candidate_end in range(nmin - 1, suffix_end): + match_len = 0 + while ( + match_len < nmax + and candidate_end - match_len >= 0 + and suffix_end - match_len >= 0 + and tokens[candidate_end - match_len] == tokens[suffix_end - match_len] + ): + match_len += 1 + if match_len >= nmin and (match_len, candidate_end) > best: + best = (match_len, candidate_end) + + if best[0] < 0: + return [0] * k, 0, 0 + match_len, candidate_end = best + valid = min(k, suffix_end - candidate_end) + continuation = tokens[candidate_end + 1 : candidate_end + 1 + valid] + return continuation + [0] * (k - valid), match_len, valid + + +@pytest.mark.parametrize("k", [7, 15]) +def test_suffix_lookup_matches_reference_and_supports_overlap(k: int) -> None: + generator = torch.Generator().manual_seed(7) + random_tokens = torch.randint(1, 8, (1500,), generator=generator).tolist() + cases = [ + random_tokens, + [5, 6, 7, 8] * 100, + list(range(1, 300)), + random_tokens[:500] + random_tokens[80:120], + ] + device = torch.device("cuda") + + for sequence in cases: + all_ids = torch.zeros((3, len(sequence) + 8), dtype=torch.int32, device=device) + all_ids[2, : len(sequence)] = torch.tensor( + sequence, dtype=torch.int32, device=device + ) + total_len = torch.zeros(3, dtype=torch.int32, device=device) + total_len[2] = len(sequence) + idx_mapping = torch.full((k,), 2, dtype=torch.int32, device=device) + eligible = torch.ones(1, dtype=torch.int32, device=device) + + actual_tokens, actual_len, actual_valid = suffix_lookup( + all_ids, + total_len, + idx_mapping, + eligible, + 1, + k, + idx_mapping_stride=k, + nmin=4, + nmax=12, + ) + expected_tokens, expected_len, expected_valid = _reference_lookup( + sequence, k, 4, 12 + ) + valid = expected_valid + assert int(actual_len[0]) == expected_len + assert int(actual_valid[0]) == expected_valid + assert actual_tokens[0, :valid].tolist() == expected_tokens[:valid] + + +def test_suffix_lookup_honors_per_request_eligibility() -> None: + device = torch.device("cuda") + all_ids = torch.tensor([[1, 2, 3, 4, 1, 2, 3, 0]], dtype=torch.int32, device=device) + total_len = torch.tensor([7], dtype=torch.int32, device=device) + idx_mapping = torch.zeros(7, dtype=torch.int32, device=device) + eligible = torch.zeros(1, dtype=torch.int32, device=device) + + tokens, match_len, valid = suffix_lookup( + all_ids, + total_len, + idx_mapping, + eligible, + 1, + 7, + idx_mapping_stride=7, + nmin=2, + nmax=4, + ) + + assert torch.count_nonzero(tokens) == 0 + assert int(match_len[0]) == 0 + assert int(valid[0]) == 0 + + +def test_fuse_decouples_seven_drafts_from_fifteen_verify_slots() -> None: + device = torch.device("cuda") + k, draft_block = 15, 7 + draft = torch.arange(100, 100 + k, dtype=torch.int64, device=device).view(1, k) + lookup = torch.arange(200, 200 + k, dtype=torch.int32, device=device).view(1, k) + use = torch.zeros((1, k), dtype=torch.int32, device=device) + hits = torch.zeros((), dtype=torch.int64, device=device) + take_flags = torch.zeros(1, dtype=torch.int32, device=device) + idx_mapping = torch.zeros(draft_block, dtype=torch.int32, device=device) + + fuse_draft( + draft, + lookup, + torch.tensor([12], dtype=torch.int32, device=device), + torch.tensor([k], dtype=torch.int32, device=device), + use, + idx_mapping, + hits, + 1, + k, + draft_block=draft_block, + idx_mapping_stride=draft_block, + nmin=4, + nstrong=6, + agree_min=0, + nmin_tail=4, + long_min=6, + take_flags=take_flags, + ) + + assert torch.equal(draft, lookup.to(torch.int64)) + assert torch.all(use == 1) + assert int(take_flags[0]) == 1 + assert int(hits) == 1 + + +def test_lookup_controller_flag_requires_lookup_and_full_q8_emission() -> None: + take = torch.tensor([1, 1, 0], dtype=torch.int32, device="cuda") + emitted = torch.tensor([8, 7, 16], dtype=torch.int32, device="cuda") + output = torch.full((3,), -1, dtype=torch.int32, device="cuda") + + _prepare_lookup_controller_flags_kernel[(1,)]( + take, + emitted, + output, + 8, + 3, + BLOCK=4, + num_warps=1, + ) + + assert output.tolist() == [1, 0, 0] + + +def test_point_mass_rewrite_preserves_sparse_cache_invariant() -> None: + device = torch.device("cuda") + num_steps, top_k, vocab_size = 15, 16, 64 + draft_logits = torch.full( + (2, num_steps, vocab_size), + -float("inf"), + dtype=torch.float32, + device=device, + ) + old_ids = ( + torch.arange(num_steps * top_k, dtype=torch.int64, device=device) + .view(num_steps, top_k) + .remainder(32) + ) + cached_ids = old_ids.repeat(2, 1, 1) + cached_scores = torch.zeros( + (2, num_steps, top_k), dtype=torch.float32, device=device + ) + draft_logits[1].scatter_(1, old_ids, cached_scores[1]) + proposed = torch.arange(40, 40 + num_steps, dtype=torch.int64, device=device) + draft_tokens = proposed.view(1, num_steps) + use = torch.ones((1, num_steps), dtype=torch.int32, device=device) + idx_mapping = torch.ones(7, dtype=torch.int32, device=device) + + _point_mass_draft_logits_kernel[(num_steps,)]( + draft_logits, + cached_ids, + cached_scores, + draft_tokens, + draft_tokens.stride(0), + use, + idx_mapping, + 7, + cached_ids.stride(0), + cached_ids.stride(1), + draft_logits.stride(0), + draft_logits.stride(1), + num_steps=num_steps, + top_k=top_k, + BLOCK_K=triton.next_power_of_2(top_k), + CACHE_SCORES=True, + num_warps=1, + ) + + for step, token in enumerate(proposed.tolist()): + finite_ids = torch.where(torch.isfinite(draft_logits[1, step]))[0].tolist() + assert finite_ids == [token] + assert float(draft_logits[1, step, token]) == 0.0 + assert int(cached_ids[1, step, 0]) == token + assert float(cached_scores[1, step, 0]) == 0.0 + assert torch.isneginf(cached_scores[1, step, 1:]).all() diff --git a/tests/v1/spec_decode/test_dflash2_ngram_assist.py b/tests/v1/spec_decode/test_dflash2_ngram_assist.py index cbf44d08f4..89ef35b2bd 100644 --- a/tests/v1/spec_decode/test_dflash2_ngram_assist.py +++ b/tests/v1/spec_decode/test_dflash2_ngram_assist.py @@ -17,6 +17,7 @@ ) from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( DFlash2Speculator, + _advance_lookup_controller, _apply_ngram_draft_kernel, ) from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( @@ -203,6 +204,94 @@ def test_ngram_assist_rejects_non_dflash_method() -> None: ) +def test_lookup_controller_requires_two_hits_and_preserves_pending_state() -> None: + state = dict( + last_want=False, + want_streak=0, + sticky_remaining=0, + long_active=False, + ) + + values = _advance_lookup_controller( + want=True, + num_reqs=1, + entry_streak=2, + sticky_steps=3, + **state, + ) + state = dict(zip(state, values)) + assert state == { + "last_want": True, + "want_streak": 1, + "sticky_remaining": 0, + "long_active": False, + } + + values = _advance_lookup_controller( + want=True, + num_reqs=1, + entry_streak=2, + sticky_steps=3, + **state, + ) + state = dict(zip(state, values)) + assert state["long_active"] is True + assert state["sticky_remaining"] == 3 + + assert _advance_lookup_controller( + want=None, + num_reqs=1, + entry_streak=2, + sticky_steps=3, + **state, + ) == tuple(state.values()) + + +def test_lookup_controller_never_coasts_a_multi_request_batch() -> None: + values = _advance_lookup_controller( + want=False, + num_reqs=2, + entry_streak=2, + sticky_steps=3, + last_want=True, + want_streak=2, + sticky_remaining=3, + long_active=True, + ) + + assert values == (False, 0, 0, False) + + +def test_lookup_controller_selects_q8_before_two_strong_hits(monkeypatch) -> None: + speculator = object.__new__(DFlash2Speculator) + speculator._lookup_enabled = True + speculator._lookup_current_req_key = (0,) + speculator._lookup_current_eligible = True + speculator._lookup_adaptive = True + speculator._lookup_cheap_context = 0 + speculator._lookup_entry_streak = 2 + speculator._lookup_sticky_steps = 3 + speculator._lookup_last_want = False + speculator._lookup_want_streak = 0 + speculator._lookup_sticky_remaining = 0 + speculator._lookup_long_active = False + speculator._lookup_last_verify_tokens = 0 + speculator._lookup_q8_rounds = 0 + speculator._lookup_q16_rounds = 0 + speculator.draft_block = 7 + speculator.num_speculative_steps = 15 + speculator.draft_max_seq_len = 32768 + wants = iter((None, True, True)) + monkeypatch.setattr(speculator, "_consume_lookup_flags", lambda: next(wants)) + monkeypatch.setattr(speculator, "_queue_lookup_flags", lambda _num_reqs: None) + + assert speculator.next_num_draft_tokens() == 7 + assert speculator.next_num_draft_tokens() == 7 + assert speculator.next_num_draft_tokens() == 15 + assert speculator._lookup_q8_rounds == 2 + assert speculator._lookup_q16_rounds == 1 + + def test_ngram_one_hot_cache_overrides_only_hit_rows() -> None: if not torch.cuda.is_available(): pytest.skip("CUDA is required for the Triton cache kernel") diff --git a/tests/v1/spec_decode/test_dflash2_structured_output.py b/tests/v1/spec_decode/test_dflash2_structured_output.py new file mode 100644 index 0000000000..fac8646da7 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash2_structured_output.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Structured-output state safety for DFlash2/speculative draft windows.""" + +from transformers import AutoTokenizer + +from vllm.config import StructuredOutputsConfig, VllmConfig +from vllm.config.model import ModelConfig +from vllm.config.speculative import SpeculativeConfig +from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.request import Request +from vllm.v1.structured_output import StructuredOutputManager + +TOKENIZER = "gpt2" +NUM_SPEC_TOKENS = 4 + + +def _make_manager_and_request(prompt_str: str = '{"a": "b"}'): + tokenizer = AutoTokenizer.from_pretrained(TOKENIZER) + prompt = tokenizer.encode(prompt_str) + config = VllmConfig( + model_config=ModelConfig(tokenizer=TOKENIZER), + structured_outputs_config=StructuredOutputsConfig(backend="xgrammar"), + speculative_config=SpeculativeConfig( + model="[ngram]", + num_speculative_tokens=NUM_SPEC_TOKENS, + ), + ) + manager = StructuredOutputManager(config) + sampling_params = SamplingParams( + structured_outputs=StructuredOutputsParams(json='{"type": "object"}'), + ) + assert sampling_params.structured_outputs is not None + sampling_params.structured_outputs._backend = "xgrammar" + sampling_params.update_from_generation_config({}, tokenizer.eos_token_id) + request = Request( + "dflash2-structured-output", + prompt_token_ids=prompt, + sampling_params=sampling_params, + pooling_params=None, + ) + manager.grammar_init(request) + assert request.structured_output_request is not None + while not request.structured_output_request._check_grammar_completion(): + pass + return tokenizer, manager, request, prompt + + +def _grammar(request: Request): + assert request.structured_output_request is not None + grammar = request.structured_output_request.grammar + assert grammar is not None + return grammar + + +def test_dflash2_bonus_row_remains_constrained_after_invalid_draft_padding(): + tokenizer, manager, request, prompt = _make_manager_and_request('{"a"') + grammar = _grammar(request) + assert grammar.accept_tokens(request.request_id, prompt) + + drafts = [tokenizer.encode(":")[0], -1, -1, -1] + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={request.request_id: drafts}, + ) + + assert bitmask is not None + assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1 + assert not (bitmask[-1] == -1).all() + assert not grammar.is_terminated() + + +def test_dflash2_post_reasoning_invalid_draft_does_not_advance_fsm(caplog): + tokenizer, manager, request, prompt = _make_manager_and_request("{") + grammar = _grammar(request) + assert grammar.accept_tokens(request.request_id, prompt) + + marker = tokenizer.encode("\n")[0] + + class MarkerReasoner: + def __init__(self, *args, **kwargs): + pass + + def is_reasoning_end(self, input_ids): + return marker in list(input_ids) + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + return marker in list(delta_ids) + + manager.reasoner_cls = MarkerReasoner + assert request.structured_output_request is not None + request.structured_output_request.reasoner = MarkerReasoner() + request.structured_output_request.reasoning_ended = False + + drafts = [tokenizer.encode(" ")[0], marker, tokenizer.encode("z")[0]] + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={request.request_id: drafts}, + ) + + assert bitmask is not None + assert (bitmask[0] == -1).all() + assert (bitmask[1] == -1).all() + assert not (bitmask[2] == -1).all() + assert not (bitmask[-1] == -1).all() + assert not grammar.is_terminated() + assert "Failed to advance FSM" not in caplog.text + + +def test_xgrammar_stops_accepting_at_termination(capfd): + tokenizer, _, request, prompt = _make_manager_and_request() + grammar = _grammar(request) + assert grammar.accept_tokens(request.request_id, prompt) + + eos = tokenizer.eos_token_id + trailing = tokenizer.encode("\n")[0] + processed_before = grammar.num_processed_tokens + + assert grammar.accept_tokens(request.request_id, [eos, trailing]) + assert grammar.is_terminated() + assert grammar.num_processed_tokens == processed_before + 1 + assert "trying to accept new token" not in capfd.readouterr().err + + processed_after = grammar.num_processed_tokens + assert grammar.accept_tokens(request.request_id, [trailing]) + assert grammar.num_processed_tokens == processed_after + + grammar.reset() + assert not grammar.is_terminated() + assert grammar.num_processed_tokens == 0 + + +def test_xgrammar_validation_rolls_back_at_termination(capfd): + tokenizer, _, request, prompt = _make_manager_and_request() + grammar = _grammar(request) + assert grammar.accept_tokens(request.request_id, prompt) + + eos = tokenizer.eos_token_id + trailing = tokenizer.encode("\n")[0] + assert grammar.validate_tokens([eos, trailing]) == [eos] + assert "trying to accept new token" not in capfd.readouterr().err + assert not grammar.matcher.is_terminated() + + assert grammar.accept_tokens(request.request_id, [eos]) + assert grammar.is_terminated() + assert grammar.validate_tokens([trailing]) == [] diff --git a/tests/v1/spec_decode/test_dflash_mrv2_config.py b/tests/v1/spec_decode/test_dflash_mrv2_config.py index 7a0292d9d2..abac531340 100644 --- a/tests/v1/spec_decode/test_dflash_mrv2_config.py +++ b/tests/v1/spec_decode/test_dflash_mrv2_config.py @@ -7,7 +7,12 @@ import pytest import torch -from vllm.config.speculative import SpeculativeConfig +from vllm.config.scheduler import SchedulerConfig +from vllm.config.speculative import ( + SpeculativeConfig, + get_dflash_model_draft_tokens, + uses_adaptive_dflash_lookup, +) from vllm.config.vllm import VllmConfig from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM from vllm.v1.spec_decode.dflash import DFlashProposer @@ -16,7 +21,10 @@ from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( get_eagle3_aux_layers_from_config, ) -from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id +from vllm.v1.worker.gpu.spec_decode.utils import ( + DraftTokensHandler, + get_parallel_drafting_token_id, +) def _config(method: str, num_speculative_tokens: int = 7) -> SpeculativeConfig: @@ -45,6 +53,108 @@ def test_mrv2_dflash_reserves_all_mask_slots() -> None: assert _config("dflash_ddtree").max_num_new_slots_for_drafting == 6 +@pytest.mark.parametrize( + ("assist", "selector_top_k", "verify_tokens", "expected"), + [ + (True, 16, 15, 7), + (False, 16, 15, 15), + (True, 0, 15, 15), + (True, 16, 7, 7), + ], +) +def test_lookup_decouples_checkpoint_and_verify_widths( + assist: bool, + selector_top_k: int, + verify_tokens: int, + expected: int, +) -> None: + config = SimpleNamespace( + method="dflash", + ngram_assist=assist, + num_speculative_tokens=verify_tokens, + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace( + dflash_config={ + "block_size": 8, + "selector_top_k": selector_top_k, + } + ) + ), + ) + + assert get_dflash_model_draft_tokens(config) == expected + assert uses_adaptive_dflash_lookup(config) is (expected < verify_tokens) + + +def test_draft_token_handler_exposes_only_selected_verify_width() -> None: + handler = object.__new__(DraftTokensHandler) + handler.req_ids = [] + handler.draft_tokens_np = None + handler.num_draft_tokens = 0 + batch = SimpleNamespace( + req_ids=["req-0"], + has_structured_output_reqs=False, + ) + proposals = torch.zeros((1, 15), dtype=torch.int64) + + handler.set_draft_tokens(batch, proposals, num_draft_tokens=7) + + assert handler.num_draft_tokens == 7 + assert handler.req_ids == ["req-0"] + with pytest.raises(ValueError, match="proposal tensor width"): + handler.set_draft_tokens(batch, proposals, num_draft_tokens=16) + + +def _adaptive_lookup_config() -> SpeculativeConfig: + config = object.__new__(SpeculativeConfig) + config.method = "dflash" + config.ngram_assist = True + config.num_speculative_tokens = 15 + config.ddtree_disable_tree_verify = False + config.disable_padded_drafter_batch = False + config.draft_model_config = SimpleNamespace( + hf_config=SimpleNamespace( + dflash_config={"block_size": 8, "selector_top_k": 16} + ), + verify_with_parallel_config=lambda *_args: None, + verify_with_model_config=lambda *_args: None, + ) + config.verify_with_parallel_config = lambda *_args: None + config.verify_with_model_config = lambda *_args: None + return config + + +def test_adaptive_lookup_disables_async_scheduling_by_default(monkeypatch) -> None: + monkeypatch.setenv("VLLM_DFLASH2_LOOKUP_ADAPTIVE", "1") + scheduler = SchedulerConfig( + max_model_len=8192, + is_encoder_decoder=False, + async_scheduling=None, + ) + + config = VllmConfig( + scheduler_config=scheduler, + speculative_config=_adaptive_lookup_config(), + ) + + assert config.scheduler_config.async_scheduling is False + + +def test_adaptive_lookup_rejects_explicit_async_scheduling(monkeypatch) -> None: + monkeypatch.setenv("VLLM_DFLASH2_LOOKUP_ADAPTIVE", "1") + scheduler = SchedulerConfig( + max_model_len=8192, + is_encoder_decoder=False, + async_scheduling=True, + ) + + with pytest.raises(ValueError, match="adaptive q8/q16 lookup verification"): + VllmConfig( + scheduler_config=scheduler, + speculative_config=_adaptive_lookup_config(), + ) + + def test_dflash_forces_v2_and_rejects_explicit_v1(monkeypatch) -> None: fake = SimpleNamespace( speculative_config=SimpleNamespace(use_dflash=lambda: True), diff --git a/tests/v1/spec_decode/test_qwen4_exp.py b/tests/v1/spec_decode/test_qwen4_exp.py new file mode 100644 index 0000000000..c430b22199 --- /dev/null +++ b/tests/v1/spec_decode/test_qwen4_exp.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Qwen4Exp MTP proposer cache topology.""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.spec_decode.qwen4_exp as qwen_proposer +from vllm.config.speculative import SpeculativeConfig +from vllm.models.qwen4_exp.common.qsa_cache import ( + circular_qsa_slot_mapping, + compressed_qsa_slot_mapping, +) +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + FullAttentionSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.spec_decode.qwen4_exp import Qwen4ExpMTPProposer + +SCHEDULER_BLOCK_SIZE = 256 +KERNEL_BLOCK_SIZE = 64 +RAW_CAPACITY = 128 +MAIN_LAYER = "draft.mtp.layers.0.self_attn.attn" +RAW_LAYER = "draft.mtp.layers.0.self_attn.indexer.raw_key_cache" +COMPRESSED_LAYER = "draft.mtp.layers.0.self_attn.indexer.compressed_key_cache" + + +class _FakeBackend: + def __init__(self, name: str) -> None: + self.name = name + + def full_cls_name(self) -> tuple[str, str]: + return (__name__, self.name) + + +class _FakeAttentionGroup: + def __init__(self, backend, layer_names, kv_cache_spec, kv_cache_group_id): + self.backend = backend + self.layer_names = list(layer_names) + self.kv_cache_spec = kv_cache_spec + self.kv_cache_group_id = kv_cache_group_id + self.builder = SimpleNamespace(kv_cache_spec=kv_cache_spec) + + def create_metadata_builders(self, vllm_config, device, kernel_block_size=None): + if kernel_block_size is not None: + self.builder.kv_cache_spec = self.kv_cache_spec.copy_with_new_block_size( + kernel_block_size + ) + + def get_metadata_builder(self): + return self.builder + + +def _make_specs(): + main_spec = FullAttentionSpec( + block_size=SCHEDULER_BLOCK_SIZE, + num_kv_heads=1, + head_size=128, + head_size_v=128, + dtype=torch.bfloat16, + ) + compressed_spec = MLAAttentionSpec( + block_size=SCHEDULER_BLOCK_SIZE, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + compress_ratio=64, + ) + raw_spec = CircularBufferSpec( + block_size=RAW_CAPACITY, + num_kv_heads=1, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + ) + return main_spec, compressed_spec, raw_spec + + +def _make_proposer_and_config( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[Qwen4ExpMTPProposer, SimpleNamespace]: + main_spec, compressed_spec, raw_spec = _make_specs() + backends = { + MAIN_LAYER: _FakeBackend("MainBackend"), + RAW_LAYER: _FakeBackend("QSAStateBackend"), + COMPRESSED_LAYER: _FakeBackend("QSAStateBackend"), + } + fake_layers = { + name: SimpleNamespace( + get_attn_backend=lambda backend=backend: backend, + num_heads=1, + ) + for name, backend in backends.items() + } + monkeypatch.setattr( + qwen_proposer, + "get_layers_from_vllm_config", + lambda *args, **kwargs: fake_layers, + ) + monkeypatch.setattr(qwen_proposer, "AttentionGroup", _FakeAttentionGroup) + + proposer = Qwen4ExpMTPProposer.__new__(Qwen4ExpMTPProposer) + proposer.vllm_config = None + proposer.draft_model_config = SimpleNamespace( + hf_text_config=SimpleNamespace(mtp_num_hidden_layers=1) + ) + proposer.device = torch.device("cpu") + proposer._draft_attn_layer_names = {MAIN_LAYER, RAW_LAYER, COMPRESSED_LAYER} + proposer.kv_cache_gid = -1 + proposer.draft_attn_groups = [] + proposer.block_size = -1 + proposer._per_group_block_tables = {} + + config = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace( + layer_names=[MAIN_LAYER, COMPRESSED_LAYER], + kv_cache_spec=UniformTypeKVCacheSpecs( + block_size=SCHEDULER_BLOCK_SIZE, + kv_cache_specs={ + MAIN_LAYER: main_spec, + COMPRESSED_LAYER: compressed_spec, + }, + ), + ), + SimpleNamespace( + layer_names=[RAW_LAYER], + kv_cache_spec=UniformTypeKVCacheSpecs( + block_size=RAW_CAPACITY, + kv_cache_specs={RAW_LAYER: raw_spec}, + ), + ), + ] + ) + return proposer, config + + +def test_initializes_and_builds_current_qwen_cache_topology( + monkeypatch: pytest.MonkeyPatch, +) -> None: + proposer, config = _make_proposer_and_config(monkeypatch) + proposer.initialize_attn_backend( + config, + kernel_block_sizes=[KERNEL_BLOCK_SIZE, RAW_CAPACITY], + ) + + assert proposer.kv_cache_gid == 0 + assert proposer.block_size == KERNEL_BLOCK_SIZE + assert [group.kv_cache_group_id for group in proposer.draft_attn_groups] == [ + 0, + 0, + 1, + ] + + logical_positions = torch.tensor([62, 63, 64, 65], dtype=torch.int64) + token_to_req = torch.zeros(4, dtype=torch.int32) + + for group in proposer.draft_attn_groups: + + def build_for_drafting(*, common_attn_metadata, draft_index, group=group): + spec = group.builder.kv_cache_spec + if isinstance(spec, CircularBufferSpec): + slots = circular_qsa_slot_mapping( + common_attn_metadata.block_table_tensor, + token_to_req, + logical_positions, + spec.block_size, + query_start_loc=common_attn_metadata.query_start_loc, + ) + elif isinstance(spec, MLAAttentionSpec): + slots = compressed_qsa_slot_mapping( + common_attn_metadata.block_table_tensor, + token_to_req, + logical_positions, + spec.storage_block_size, + spec.compress_ratio, + ) + else: + slots = common_attn_metadata.slot_mapping + return SimpleNamespace( + common=common_attn_metadata, + draft_index=draft_index, + slot_mapping=slots, + ) + + group.builder.build_for_drafting = build_for_drafting + + main_table = torch.tensor([[10, 11, 12, 13]], dtype=torch.int32) + raw_table = torch.tensor([[20]], dtype=torch.int32) + main_slots = torch.tensor([702, 703, 704, 705], dtype=torch.int64) + proposer.set_per_group_block_table(1, raw_table) + common = SimpleNamespace( + num_reqs=1, + query_start_loc=torch.tensor([0, 4], dtype=torch.int32), + block_table_tensor=main_table, + slot_mapping=main_slots, + ) + + _, per_layer = proposer.build_per_group_and_layer_attn_metadata( + common, + draft_index=2, + ) + + assert torch.equal(per_layer[MAIN_LAYER].slot_mapping, main_slots) + assert torch.equal( + per_layer[COMPRESSED_LAYER].slot_mapping, + torch.tensor([-1, 10, -1, -1], dtype=torch.int64), + ) + assert torch.equal( + per_layer[RAW_LAYER].slot_mapping, + torch.tensor([2622, 2623, 2624, 2625], dtype=torch.int64), + ) + + +def test_rejects_multiple_mtp_layers(monkeypatch: pytest.MonkeyPatch) -> None: + proposer, config = _make_proposer_and_config(monkeypatch) + proposer.draft_model_config.hf_text_config.mtp_num_hidden_layers = 2 + + with pytest.raises(NotImplementedError, match="only supports one MTP layer"): + proposer.initialize_attn_backend( + config, + kernel_block_sizes=[KERNEL_BLOCK_SIZE, RAW_CAPACITY], + ) + + +def test_qwen_proposer_hidden_and_config_contract() -> None: + proposer = Qwen4ExpMTPProposer.__new__(Qwen4ExpMTPProposer) + proposer.draft_model_config = SimpleNamespace( + hf_config=SimpleNamespace(hc_mult=4), + get_hidden_size=lambda: 1024, + ) + config = SimpleNamespace( + method="mtp", + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace(model_type="qwen4_exp_mtp") + ), + ) + + assert proposer._get_hidden_size() == 4096 + assert proposer.model_returns_tuple() + assert SpeculativeConfig.use_qwen4_exp_mtp(config) diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 83bb10599e..6ca495528b 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -132,9 +132,11 @@ def _assert_distribution_match( @pytest.mark.parametrize("num_speculative_steps", [3, 7]) @pytest.mark.parametrize("top_p", [1.0, 0.95]) +@pytest.mark.parametrize("temperature", [0.6, 1.0]) def test_dflash2_sparse_topk_matches_dense_rejection( num_speculative_steps: int, top_p: float, + temperature: float, ): """Compact p/q support must preserve the dense DFlash2 decision path.""" torch.manual_seed(20260823) @@ -174,7 +176,11 @@ def test_dflash2_sparse_topk_matches_dense_rejection( top_p_rows = torch.full((num_logits,), top_p, dtype=torch.float32, device=device) from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p - processed_target = apply_top_k_top_p(target_dense.clone(), top_k_tensor, top_p_rows) + processed_target = apply_top_k_top_p( + target_dense.clone() / temperature, + top_k_tensor, + top_p_rows, + ) draft_topk_ids = torch.stack( [ @@ -187,12 +193,15 @@ def test_dflash2_sparse_topk_matches_dense_rejection( for _ in range(num_reqs) ] ) - draft_topk_logits = torch.randn( - num_reqs, - num_speculative_steps, - draft_top_k, - dtype=torch.float32, - device=device, + draft_topk_logits = ( + torch.randn( + num_reqs, + num_speculative_steps, + draft_top_k, + dtype=torch.float32, + device=device, + ) + / temperature ) draft_dense = torch.full( (num_reqs, num_speculative_steps, VOCAB_SIZE), @@ -216,7 +225,9 @@ def test_dflash2_sparse_topk_matches_dense_rejection( rows_per_req, dtype=torch.int32, device=device ).repeat(num_reqs) pos = torch.arange(num_logits, dtype=torch.int64, device=device) + 8192 - temperature = torch.ones(num_reqs, dtype=torch.float32, device=device) + temperature_per_req = torch.full( + (num_reqs,), temperature, dtype=torch.float32, device=device + ) top_p_per_req = torch.full((num_reqs,), top_p, dtype=torch.float32, device=device) seeds = torch.arange(101, 101 + num_reqs, dtype=torch.int64, device=device) @@ -229,7 +240,7 @@ def test_dflash2_sparse_topk_matches_dense_rejection( idx_mapping, expanded_idx_mapping, expanded_local_pos, - temperature, + temperature_per_req, seeds, num_speculative_steps, ) @@ -242,7 +253,7 @@ def test_dflash2_sparse_topk_matches_dense_rejection( cu_num_logits, pos, idx_mapping, - temperature, + temperature_per_req, top_p_per_req, seeds, num_speculative_steps, diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py index 861e919c10..bbb0dd1d89 100644 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ b/tests/v1/structured_output/test_reasoning_structured_output.py @@ -189,7 +189,7 @@ def test_should_advance_reasoning_not_ended( # Should return False since reasoning hasn't ended assert result is False - def test_should_advance_reasoning_just_ended( + def test_should_advance_reasoning_just_ended_records_boundary( self, manager_with_reasoner, mock_request_with_structured_output, @@ -204,16 +204,23 @@ def test_should_advance_reasoning_just_ended( structured_req = mock_request_with_structured_output.structured_output_request structured_req.reasoner = reasoner + new_token_ids = [6, 7, 8] result = manager_with_reasoner.should_advance( - mock_request_with_structured_output + mock_request_with_structured_output, + new_token_ids=new_token_ids, ) - # Should set reasoning_ended to True but return False for this step + # The scheduler must advance only the suffix after the marker in the + # same step; deferring loses that suffix under speculative decoding. assert ( mock_request_with_structured_output.structured_output_request.reasoning_ended is True ) - assert result is False + assert result is True + assert ( + mock_request_with_structured_output.structured_output_request.reasoning_end_token_index + == 5 + ) def test_should_advance_reasoning_just_ended_with_spec_decode_structural_tag( self, @@ -241,6 +248,46 @@ def test_should_advance_reasoning_just_ended_with_spec_decode_structural_tag( assert structured_req.reasoning_ended is True assert result is True + def test_should_advance_uses_appended_tokens_not_stale_placeholders( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_ended = False + mock_request_with_structured_output.num_output_placeholders = 99 + + reasoner = MockReasoner(tokenizer=Mock()) + reasoner.is_reasoning_end_streaming.side_effect = ( + lambda _all_ids, delta_ids: 8 in list(delta_ids) + ) + structured_req.reasoner = reasoner + + assert manager_with_reasoner.should_advance( + mock_request_with_structured_output, + new_token_ids=[7, 8], + ) + assert structured_req.reasoning_end_token_index == 7 + + def test_trim_reasoning_for_advance_keeps_only_post_marker_suffix( + self, + manager_with_reasoner, + mock_request_with_structured_output, + ): + structured_req = mock_request_with_structured_output.structured_output_request + structured_req.reasoning_end_token_index = 6 + + assert manager_with_reasoner.trim_reasoning_for_advance( + mock_request_with_structured_output, + [6, 7, 8], + ) == [8] + + mock_request_with_structured_output.all_token_ids.extend([9, 10]) + assert manager_with_reasoner.trim_reasoning_for_advance( + mock_request_with_structured_output, + [9, 10], + ) == [9, 10] + def test_should_advance_reasoning_already_ended( self, manager_with_reasoner, diff --git a/tests/v1/worker/test_glm5_compressed_kv_cache.py b/tests/v1/worker/test_glm5_compressed_kv_cache.py new file mode 100644 index 0000000000..acd3878c10 --- /dev/null +++ b/tests/v1/worker/test_glm5_compressed_kv_cache.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + MLAAttentionSpec, +) +from vllm.v1.worker.gpu_model_runner import GPUModelRunner +from vllm.v1.worker.utils import AttentionGroup + + +class _IndexerBackend: + @staticmethod + def get_kv_cache_shape( + num_blocks, + block_size, + num_kv_heads, + head_size, + cache_dtype_str="auto", + ): + del cache_dtype_str + assert num_kv_heads == 1 + return (num_blocks, block_size, head_size) + + @staticmethod + def get_kv_cache_stride_order(): + return (0, 1, 2) + + +def test_glm53_indexer_cache_uses_physical_pool_pages(): + spec = MLAAttentionSpec( + block_size=1152, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + compress_ratio=4, + ) + layer_name = "model.layers.3.self_attn.indexer.k_cache" + group = AttentionGroup( + backend=_IndexerBackend, + layer_names=[layer_name], + kv_cache_spec=spec, + kv_cache_group_id=0, + ) + + runner = object.__new__(GPUModelRunner) + runner.kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec([layer_name], spec)], + ) + runner.attn_groups = [[group]] + runner.runner_only_attn_layers = set() + runner.cache_config = SimpleNamespace(cache_dtype="auto") + + raw = torch.empty(spec.page_size_bytes, dtype=torch.int8) + cache = runner._reshape_kv_cache_tensors({layer_name: raw}, [64])[layer_name] + + # 1152 scheduler tokens / kpool=4 = 288 stored entries, split into the + # 32-entry pages accepted by the paged MQA kernel. + assert cache.shape == (9, 32, 132) + assert cache.untyped_storage().data_ptr() == raw.untyped_storage().data_ptr() diff --git a/tests/v1/worker/test_gpu_cudagraph_memory_reserve.py b/tests/v1/worker/test_gpu_cudagraph_memory_reserve.py new file mode 100644 index 0000000000..d2471c1d9d --- /dev/null +++ b/tests/v1/worker/test_gpu_cudagraph_memory_reserve.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.worker.gpu.cudagraph_utils import ( + get_explicit_cudagraph_memory_reserve, +) + +pytestmark = pytest.mark.cpu_test + + +def test_explicit_cudagraph_memory_reserve(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_V2_CUDAGRAPH_MEM_MIB", "1536.5") + + assert get_explicit_cudagraph_memory_reserve(CUDAGraphMode.FULL) == int( + 1536.5 * 1024 * 1024 + ) + assert get_explicit_cudagraph_memory_reserve(CUDAGraphMode.NONE) == 0 + + +def test_explicit_cudagraph_memory_reserve_rejects_negative( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv("VLLM_V2_CUDAGRAPH_MEM_MIB", "-1") + + with pytest.raises(ValueError, match="must be non-negative"): + get_explicit_cudagraph_memory_reserve(CUDAGraphMode.FULL) diff --git a/tests/v1/worker/test_gpu_model_runner_pp.py b/tests/v1/worker/test_gpu_model_runner_pp.py new file mode 100644 index 0000000000..9e03c6fd0b --- /dev/null +++ b/tests/v1/worker/test_gpu_model_runner_pp.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np +import torch + +from vllm.sequence import IntermediateTensors +from vllm.v1.worker.gpu_model_runner import _select_dummy_sample_hidden_states + + +def test_pp_dummy_intermediates_are_not_sampled() -> None: + intermediate = IntermediateTensors( + {"hidden_states": torch.empty((8, 16), dtype=torch.float16)} + ) + + result = _select_dummy_sample_hidden_states( + intermediate, np.array([3, 5]), torch.device("cpu") + ) + + assert result is None + + +def test_last_pp_rank_selects_final_scheduled_tokens() -> None: + hidden_states = torch.arange(8 * 2).view(8, 2) + + result = _select_dummy_sample_hidden_states( + hidden_states, np.array([3, 5]), torch.device("cpu") + ) + + assert result is not None + torch.testing.assert_close(result, hidden_states[[2, 7]]) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_greedy.py b/tests/v1/worker/test_gpu_model_runner_v2_greedy.py new file mode 100644 index 0000000000..92f8b2d9bc --- /dev/null +++ b/tests/v1/worker/test_gpu_model_runner_v2_greedy.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +import vllm.v1.worker.gpu.sample.sampler as sampler_module +from vllm.v1.worker.gpu import model_runner as model_runner_module +from vllm.v1.worker.gpu.model_runner import GPUModelRunner +from vllm.v1.worker.gpu.sample.sampler import Sampler + + +def _sampler() -> Sampler: + sampler = object.__new__(Sampler) + sampler.compute_nans = False + sampler.sampling_states = SimpleNamespace( + temperature=SimpleNamespace(np=np.array([0.0, 0.0], dtype=np.float32)), + max_num_logprobs=lambda _indices: -1, + ) + sampler.logprob_token_ids_state = SimpleNamespace( + max_num_token_ids=lambda _indices: 0 + ) + sampler.logit_bias_state = SimpleNamespace(use_logit_bias=np.array([False, False])) + sampler.penalties_state = SimpleNamespace(use_penalty=np.array([False, False])) + sampler.bad_words_state = SimpleNamespace( + num_bad_words=SimpleNamespace(np=np.array([0, 0], dtype=np.int32)) + ) + return sampler + + +def test_sm70_v2_greedy_fastpath_accepts_plain_temperature_zero( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(sampler_module.envs, "VLLM_SM70_GREEDY_TOKEN_FASTPATH", True) + sampler = _sampler() + input_batch = SimpleNamespace(idx_mapping_np=np.array([1], dtype=np.int32)) + + assert sampler.can_use_sm70_greedy_token_fastpath(input_batch) + + +def test_sm70_v2_decode_uses_model_top_tokens( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Model: + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert hidden_states.shape == (1, 4) + return torch.tensor([42], dtype=torch.int64) + + def compute_logits(self, _hidden_states: torch.Tensor) -> torch.Tensor: + raise AssertionError("full logits must not run on the greedy fastpath") + + runner = object.__new__(GPUModelRunner) + runner.model = Model() + runner.lora_config = None + runner.sampler = SimpleNamespace( + can_use_sm70_greedy_token_fastpath=lambda _input_batch: True + ) + runner.device = torch.device("cuda") + runner.req_states = SimpleNamespace(prefill_len=SimpleNamespace(gpu=None)) + input_batch = SimpleNamespace( + logits_indices=torch.tensor([0]), + num_draft_tokens=0, + num_reqs=1, + num_tokens=1, + is_prefilling_np=np.array([False]), + seq_lens=torch.tensor([17], dtype=torch.int32), + cu_num_logits=None, + idx_mapping=None, + ) + monkeypatch.setattr( + model_runner_module.current_platform, + "is_device_capability", + lambda _capability: True, + ) + monkeypatch.setattr( + model_runner_module, + "get_num_sampled_and_rejected", + lambda *_args: ( + torch.ones(1, dtype=torch.int32), + torch.zeros(1, dtype=torch.int32), + ), + ) + + output, num_sampled, num_rejected = GPUModelRunner.sample( + runner, torch.zeros(1, 4), input_batch, None + ) + + assert output.sampled_token_ids.tolist() == [[42]] + assert num_sampled.tolist() == [1] + assert num_rejected.tolist() == [0] + + runner.lora_config = object() + with pytest.raises(AssertionError, match="full logits must not run"): + GPUModelRunner.sample(runner, torch.zeros(1, 4), input_batch, None) + + +@pytest.mark.parametrize( + "blocker", + [ + "disabled", + "nan_count", + "random", + "logprobs", + "logprob_token_ids", + "logit_bias", + "penalty", + "bad_words", + ], +) +def test_sm70_v2_greedy_fastpath_rejects_non_equivalent_sampling( + monkeypatch: pytest.MonkeyPatch, blocker: str +) -> None: + monkeypatch.setattr( + sampler_module.envs, + "VLLM_SM70_GREEDY_TOKEN_FASTPATH", + blocker != "disabled", + ) + sampler = _sampler() + input_batch = SimpleNamespace(idx_mapping_np=np.array([1], dtype=np.int32)) + + if blocker == "nan_count": + sampler.compute_nans = True + elif blocker == "random": + sampler.sampling_states.temperature.np[1] = 0.5 + elif blocker == "logprobs": + sampler.sampling_states.max_num_logprobs = lambda _indices: 1 + elif blocker == "logprob_token_ids": + sampler.logprob_token_ids_state.max_num_token_ids = lambda _indices: 1 + elif blocker == "logit_bias": + sampler.logit_bias_state.use_logit_bias[1] = True + elif blocker == "penalty": + sampler.penalties_state.use_penalty[1] = True + elif blocker == "bad_words": + sampler.bad_words_state.num_bad_words.np[1] = 1 + + assert not sampler.can_use_sm70_greedy_token_fastpath(input_batch) diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index b4fcb8bd0a..4beda08b46 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -7,7 +7,12 @@ from vllm.config.speculative import SpeculativeConfig from vllm.config.vllm import VllmConfig -from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + FullAttentionSpec, + MambaSpec, + UniformTypeKVCacheSpecs, +) from vllm.v1.worker.gpu.warmup import _reserved_block_count BLOCK_SIZE = 16 @@ -73,10 +78,26 @@ def _mamba_spec(mode: str) -> MambaSpec: ) +def _circular_spec() -> CircularBufferSpec: + return CircularBufferSpec( + block_size=4, + num_kv_heads=1, + head_size=8, + dtype=torch.float16, + ) + + @pytest.mark.parametrize( ("spec", "expected"), [ (_full_attention_spec(), 2), + (_circular_spec(), 1), + ( + UniformTypeKVCacheSpecs( + block_size=4, kv_cache_specs={"compressor_state": _circular_spec()} + ), + 1, + ), (_mamba_spec("align"), 8), (_mamba_spec("none"), 9), (_mamba_spec("all"), 9), diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index 332111b7f4..e7ba9e3cd3 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -159,3 +159,81 @@ def test_zero_block_ids_multiple_interleaved_pools() -> None: index for index in range(cells.shape[0]) if bool((cells[index] == 0).all()) } assert zeroed == expected + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_zero_block_ids_nonuniform_page_sizes() -> None: + device = torch.device("cuda") + page_sizes = [8, 16] + buffers = [ + torch.ones(N_BLOCKS * page_size, dtype=torch.int32, device=device) + for page_size in page_sizes + ] + groups = [] + static_forward_context = {} + for index, (page_size, buffer) in enumerate(zip(page_sizes, buffers)): + spec = MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.int32,), + page_size_padded=page_size * 4, + ) + layer_name = f"mamba.nonuniform.{index}" + groups.append( + SimpleNamespace( + kv_cache_spec=spec, + kv_cache_group_id=index, + layer_names=[layer_name], + backend=None, + ) + ) + static_forward_context[layer_name] = SimpleNamespace(kv_cache=[buffer]) + + zeroer = KVBlockZeroer(device, pin_memory=False) + zeroer.init_meta( + groups, + kernel_block_sizes=[16, 16], + cache_dtype="auto", + runner_only_attn_layers=set(), + static_forward_context=static_forward_context, + ) + + target = 2 + zeroer.zero_block_ids([target]) + torch.accelerator.synchronize() + + for page_size, buffer in zip(page_sizes, buffers): + pages = buffer.view(N_BLOCKS, page_size) + zeroed = {index for index in range(N_BLOCKS) if bool((pages[index] == 0).all())} + assert zeroed == {target} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_zeroer_supports_heterogeneous_page_sizes() -> None: + device = torch.device("cuda") + small = torch.ones((3, 8), dtype=torch.int32, device=device) + large = torch.ones((3, 20), dtype=torch.int32, device=device) + zeroer = KVBlockZeroer(device, pin_memory=False) + zeroer._id_cap = 8 + zeroer._ids_pinned = torch.empty(8, dtype=torch.int64) + zeroer._ids_gpu = torch.empty(8, dtype=torch.int64, device=device) + zeroer._meta = ( + torch.tensor( + [small.data_ptr(), large.data_ptr()], dtype=torch.uint64, device=device + ), + torch.tensor([8, 20], dtype=torch.int64, device=device), + torch.tensor([8, 20], dtype=torch.int64, device=device), + 3, + 8, + 2, + ) + + zeroer.zero_block_ids([1]) + torch.accelerator.synchronize() + + torch.testing.assert_close(small[0], torch.ones_like(small[0])) + torch.testing.assert_close(small[1], torch.zeros_like(small[1])) + torch.testing.assert_close(small[2], torch.ones_like(small[2])) + torch.testing.assert_close(large[0], torch.ones_like(large[0])) + torch.testing.assert_close(large[1], torch.zeros_like(large[1])) + torch.testing.assert_close(large[2], torch.ones_like(large[2])) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index fc2ecea746..7c1d40500d 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -13,6 +13,7 @@ get_conv_copy_spec, get_temporal_copy_spec, ) +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpec from vllm.v1.worker import mamba_utils as worker_mamba_utils @@ -469,12 +470,103 @@ def make_buffer(n, dtype): return MambaSpecDecodeGPUContext.create( max_num_reqs=cfg.max_num_reqs, kv_cache_config=kv_cache_config, - num_state_types=2, + mamba_state_copy_funcs=_COPY_FUNCS, device=device, make_buffer=make_buffer, ) +def test_mamba_context_supports_heterogeneous_state_groups(): + device = torch.device("cpu") + block_size = 16 + num_speculative_blocks = 4 + gdn_spec = MambaSpec( + block_size=block_size, + shapes=((4, 8), (16,)), + dtypes=(torch.float16, torch.float32), + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + mamba_cache_mode="align", + num_speculative_blocks=num_speculative_blocks, + ) + ple_spec = MambaSpec( + block_size=block_size, + shapes=((3, 8),), + dtypes=(torch.float16,), + mamba_type=MambaAttentionBackendEnum.SHORT_CONV, + mamba_cache_mode="align", + num_speculative_blocks=num_speculative_blocks, + tp_replicated=True, + ) + kv_cache_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=["gdn_0", "gdn_1"], kv_cache_spec=gdn_spec), + KVCacheGroupSpec(layer_names=["ple"], kv_cache_spec=ple_spec), + ], + ) + copy_funcs_by_type = { + MambaAttentionBackendEnum.GDN_ATTN: ( + get_conv_copy_spec, + get_temporal_copy_spec, + ), + MambaAttentionBackendEnum.SHORT_CONV: (get_conv_copy_spec,), + } + + def make_buffer(n, dtype): + return _MockCpuGpuBuffer(n, dtype, device) + + ctx = MambaSpecDecodeGPUContext.create( + max_num_reqs=2, + kv_cache_config=kv_cache_config, + mamba_state_copy_funcs=copy_funcs_by_type, + device=device, + make_buffer=make_buffer, + ) + copy_bufs = MambaCopyBuffers.create( + max_num_reqs=2, + kv_cache_config=kv_cache_config, + copy_funcs=copy_funcs_by_type, + make_buffer=make_buffer, + ) + + assert ctx.num_layers == 3 + assert ctx.num_states == 5 + assert ctx.state_base_addrs.numel() == 5 + assert copy_bufs.src_ptrs.cpu.numel() == 10 + + def attention(*states: torch.Tensor) -> MagicMock: + mock = MagicMock() + mock.kv_cache = list(states) + return mock + + forward_context = { + "gdn_0": attention( + torch.zeros(4, 4, 8, dtype=torch.float16), + torch.zeros(4, 16, dtype=torch.float32), + ), + "gdn_1": attention( + torch.zeros(4, 4, 8, dtype=torch.float16), + torch.zeros(4, 16, dtype=torch.float32), + ), + "ple": attention(torch.zeros(4, 3, 8, dtype=torch.float16)), + } + block_tables = [ + torch.zeros(2, 8, dtype=torch.int32), + torch.zeros(2, 8, dtype=torch.int32), + ] + ctx.initialize_from_forward_context( + kv_cache_config, + forward_context, + copy_funcs_by_type, + block_tables, + ) + + assert ctx.state_group_indices.tolist() == [0, 0, 0, 0, 1] + assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 3] + assert ctx.state_inner_sizes.tolist() == [8, 16, 8, 16, 8] + + def _run_gpu_postprocess( gpu_ctx: MambaSpecDecodeGPUContext, *, diff --git a/tests/v1/worker/test_model_state_init.py b/tests/v1/worker/test_model_state_init.py new file mode 100644 index 0000000000..44200b5806 --- /dev/null +++ b/tests/v1/worker/test_model_state_init.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch +from torch import nn + +from vllm.v1.worker.gpu.model_states import init_model_state + + +def test_model_can_select_custom_model_state() -> None: + captured: dict[str, object] = {} + + class CustomModelState: + def __init__(self, vllm_config, model, encoder_cache, device) -> None: + captured.update( + vllm_config=vllm_config, + model=model, + encoder_cache=encoder_cache, + device=device, + ) + + class CustomModel(nn.Module): + @staticmethod + def get_model_state_cls(): + return CustomModelState + + model = CustomModel() + vllm_config = SimpleNamespace() + device = torch.device("cpu") + + model_state = init_model_state(vllm_config, model, None, device) + + assert isinstance(model_state, CustomModelState) + assert captured == { + "vllm_config": vllm_config, + "model": model, + "encoder_cache": None, + "device": device, + } diff --git a/tests/v1/worker/test_qwen4_exp_ngram.py b/tests/v1/worker/test_qwen4_exp_ngram.py new file mode 100644 index 0000000000..6c2459c897 --- /dev/null +++ b/tests/v1/worker/test_qwen4_exp_ngram.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import torch + +from vllm.v1.worker.gpu_model_runner import GPUModelRunner + + +class _Buffer: + def __init__(self, shape: tuple[int, ...]) -> None: + self.np = np.zeros(shape, dtype=np.int32) + self.gpu = torch.zeros(shape, dtype=torch.int32) + + def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: + if n is None: + self.gpu.copy_(torch.from_numpy(self.np)) + else: + self.gpu[:n].copy_(torch.from_numpy(self.np[:n])) + return self.gpu + + +def _runner() -> GPUModelRunner: + runner = object.__new__(GPUModelRunner) + runner.uses_ngram_embedding = True + runner.ngram_context_len = 2 + runner.ngram_eos_token_id = 99 + runner.enable_prompt_embeds = False + runner._sm70_async_staged_input_prep_active = False + runner.ngram_context = _Buffer((4, 2)) + runner.query_start_loc = _Buffer((5,)) + runner.input_batch = SimpleNamespace( + num_computed_tokens_cpu=np.array([3, 1, 0, 0], dtype=np.int32), + token_ids_cpu=np.array( + [ + [1, 2, 3, 4], + [20, 21, 22, 23], + [0, 0, 0, 0], + [0, 0, 0, 0], + ], + dtype=np.int32, + ), + is_token_ids=np.ones((4, 4), dtype=bool), + ) + return runner + + +def test_v1_runner_prepares_committed_ngram_context() -> None: + context = _runner()._prepare_ngram_context(num_reqs=2, num_reqs_padded=3) + + torch.testing.assert_close( + context, + torch.tensor([[2, 3], [99, 20], [99, 99]], dtype=torch.int32), + ) + + +def test_v1_runner_uses_stable_dummy_ngram_buffers() -> None: + runner = _runner() + model_kwargs: dict[str, Any] = {} + + runner._maybe_add_ngram_kwargs( + model_kwargs, + num_reqs=2, + num_reqs_padded=3, + is_first_rank=True, + is_encoder_decoder=False, + use_dummy_context=True, + num_scheduled_tokens=[2, 1], + ) + + torch.testing.assert_close( + model_kwargs["query_start_loc"], + torch.tensor([0, 2, 3, 3], dtype=torch.int32), + ) + torch.testing.assert_close( + model_kwargs["ngram_context"], + torch.full((3, 2), 99, dtype=torch.int32), + ) diff --git a/tests/v1/worker/test_qwen4_exp_v2.py b/tests/v1/worker/test_qwen4_exp_v2.py new file mode 100644 index 0000000000..954779157f --- /dev/null +++ b/tests/v1/worker/test_qwen4_exp_v2.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.worker.gpu import model_runner as mrv2 +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.spec_decode.eagle import speculator as eagle_speculator + + +def test_qwen4_exp_mtp_v2_unpacks_logits_and_feedback_hidden_states( + monkeypatch: pytest.MonkeyPatch, +) -> None: + speculator = eagle_speculator.EagleSpeculator.__new__( + eagle_speculator.EagleSpeculator + ) + speculator.device = torch.device("cpu") + speculator.vllm_config = SimpleNamespace() + speculator.supports_mm_inputs = False + speculator.input_buffers = SimpleNamespace( + input_ids=torch.zeros(3, dtype=torch.int64), + positions=torch.arange(3, dtype=torch.int64), + ) + speculator.hidden_states = torch.zeros(3, 16) + + logits_hidden = torch.ones(3, 4) + feedback_hidden = torch.ones(3, 16) + speculator.model = lambda **_kwargs: (logits_hidden, feedback_hidden) + monkeypatch.setattr( + eagle_speculator, + "set_forward_context", + lambda *_args, **_kwargs: nullcontext(), + ) + + actual_logits_hidden, actual_feedback_hidden = speculator.run_model( + num_tokens=3, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=None, + ) + + assert actual_logits_hidden is logits_hidden + assert actual_feedback_hidden is feedback_hidden + + +def test_qwen4_exp_mtp_v2_uses_local_argmax_without_full_logits() -> None: + speculator = eagle_speculator.EagleSpeculator.__new__( + eagle_speculator.EagleSpeculator + ) + speculator.use_local_argmax_reduction = True + + class DraftModel: + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert hidden_states.shape == (2, 4) + return torch.tensor([7, 11]) + + def compute_logits(self, _hidden_states: torch.Tensor) -> torch.Tensor: + raise AssertionError("local argmax must not materialize full logits") + + speculator.model = DraftModel() + top_tokens = speculator._sample_draft( + hidden_states=torch.zeros(2, 4), + idx_mapping=torch.arange(2, dtype=torch.int32), + pos=torch.arange(2), + draft_step=torch.tensor(0), + draft_logits=None, + ) + + torch.testing.assert_close(top_tokens, torch.tensor([7, 11])) + + +def test_qwen4_exp_mtp_v2_reuses_step_zero_qsa_indices() -> None: + calls: list[tuple[str, object]] = [] + + class MTPBackbone: + def set_skip_topk(self, skip: bool) -> None: + calls.append(("skip", skip)) + + def compact_topk_indices(self, row_indices: torch.Tensor) -> None: + calls.append(("compact", row_indices.tolist())) + + speculator = eagle_speculator.EagleSpeculator.__new__( + eagle_speculator.EagleSpeculator + ) + speculator.share_mtp_topk_indices = True + speculator.num_speculative_steps = 3 + speculator.last_token_indices = torch.tensor([4, 9, 12]) + speculator.model = SimpleNamespace(model=MTPBackbone()) + + speculator._mtp_prefill_begin() + speculator._mtp_prefill_end(num_reqs=2) + speculator._mtp_decode_begin() + speculator._mtp_decode_end() + + assert calls == [ + ("skip", False), + ("compact", [4, 9]), + ("skip", True), + ("skip", False), + ] + + +def test_qsa_circular_group_uses_one_block_and_custom_slot_mapping( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner) + runner.max_model_len = 262144 + runner.is_encoder_decoder = False + runner.dcp_size = 1 + runner.dcp_rank = 0 + runner.cp_interleave = 1 + runner.cache_config = SimpleNamespace(enable_prefix_caching=False) + runner.vllm_config = SimpleNamespace() + runner.max_num_reqs = 1 + runner.max_num_tokens = 2 + runner.device = torch.device("cpu") + + circular_spec = CircularBufferSpec( + block_size=8, + num_kv_heads=1, + head_size=128, + dtype=torch.float16, + ) + attention_spec = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.float16, + ) + compressed_spec = MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.float16, + ) + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + layer_names=["compressor_state"], + kv_cache_spec=UniformTypeKVCacheSpecs( + block_size=8, kv_cache_specs={"compressor_state": circular_spec} + ), + ), + KVCacheGroupSpec( + layer_names=["attention", "compressed"], + kv_cache_spec=UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={ + "attention": attention_spec, + "compressed": compressed_spec, + }, + ), + ), + ], + ) + + monkeypatch.setattr( + mrv2, + "init_attn_backend", + lambda *args: ([], SimpleNamespace(), [8, 16]), + ) + captured: dict[str, object] = {} + + class BlockTablesCaptured(Exception): + pass + + def capture_block_tables(**kwargs): + captured.update(kwargs) + raise BlockTablesCaptured + + monkeypatch.setattr(mrv2, "BlockTables", capture_block_tables) + + with pytest.raises(BlockTablesCaptured): + runner.initialize_kv_cache(kv_cache_config) + + assert captured["max_num_blocks_per_group"] == [1, 16384] + assert captured["slot_mapping_enabled"] == [False, True] + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), + reason="requires an exact SM70 CUDA device", +) +def test_qsa_circular_group_emits_no_generic_slots_on_sm70() -> None: + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[8, 262144], + max_num_reqs=1, + max_num_batched_tokens=4, + max_num_blocks_per_group=[1, 1], + device=device, + kernel_block_sizes=[8, 262144], + slot_mapping_enabled=[False, True], + ) + block_tables.append_block_ids( + req_index=0, + new_block_ids=([7], [12]), + overwrite=True, + ) + block_tables.apply_staged_writes() + + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping=torch.tensor([0], dtype=torch.int32, device=device), + query_start_loc=torch.tensor([0, 2], dtype=torch.int32, device=device), + positions=torch.tensor([153797, 165757], dtype=torch.int64, device=device), + num_tokens_padded=2, + ) + torch.accelerator.synchronize() + + assert slot_mappings[0].tolist() == [-1, -1] + assert slot_mappings[1].tolist() == [ + 12 * 262144 + 153797, + 12 * 262144 + 165757, + ] diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 365b2f5bb7..b71d64a677 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -80,7 +80,12 @@ class ForbiddenImport: allowed_pattern=re.compile( "from vllm.triton_utils import (triton|tl|tl, triton)" ), - allowed_files={"vllm/triton_utils/importing.py"}, + allowed_files={ + "vllm/triton_utils/importing.py", + # AOT inspection needs compiler/backend modules that the vLLM + # runtime shim deliberately does not export. + "benchmarks/kernels/deepseek_v4_sm70_aot_check.py", + }, ), } diff --git a/tools/pre_commit/check_torch_cuda.py b/tools/pre_commit/check_torch_cuda.py index bfbb36ffbf..2511f241ba 100644 --- a/tools/pre_commit/check_torch_cuda.py +++ b/tools/pre_commit/check_torch_cuda.py @@ -19,7 +19,17 @@ "vllm/platforms/", "vllm/device_allocator/", "vllm/distributed/weight_transfer/ipc_engine.py", + # Standalone SM70/DSpark probes intentionally exercise CUDA-only APIs. + # Keep the runtime package on torch.accelerator while admitting these + # hardware-specific benchmark and operator-test namespaces as baseline. + "benchmarks/benchmark_sm70_", + "benchmarks/kernels/benchmark_flash_v100_", + "benchmarks/kernels/benchmark_sm70_", + "benchmarks/kernels/benchmark_dspark_", + "benchmarks/kernels/benchmark_deepseek_v4_sm70_", "tests/distributed/test_packed_tensor.py", + "tests/kernels/test_sm70_", + "tests/kernels/quantization/test_sm70_", } diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 8d8d0b188b..f2bea0da4b 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os from enum import IntEnum from typing import TYPE_CHECKING, Literal @@ -2963,13 +2964,30 @@ def get_max_shared_memory_per_block_device_attribute(device: int) -> int: # custom ar +def _maybe_load_sm70_custom_ar_library() -> None: + library_path = os.getenv("VLLM_SM70_CUSTOM_AR_LIBRARY") + if library_path is not None: + torch.ops.load_library(library_path) + + +_maybe_load_sm70_custom_ar_library() + + +def _custom_ar_op(name: str): + """Prefer an opt-in SM70 sidecar while retaining main-op fallbacks.""" + sidecar = torch.ops._C_custom_ar_flashnext + if hasattr(sidecar, name): + return getattr(sidecar, name) + return getattr(torch.ops._C_custom_ar, name) + + def init_custom_ar( ipc_tensors: list[torch.Tensor], rank_data: torch.Tensor, rank: int, fully_connected: bool, ) -> int: - return torch.ops._C_custom_ar.init_custom_ar( + return _custom_ar_op("init_custom_ar")( ipc_tensors, rank_data, rank, fully_connected ) @@ -2981,7 +2999,7 @@ def all_reduce( reg_buffer: int, reg_buffer_sz_bytes: int, ) -> None: - torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) + _custom_ar_op("all_reduce")(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) def sm70_tp2_all_reduce_gemma_rms_norm( @@ -2995,7 +3013,7 @@ def sm70_tp2_all_reduce_gemma_rms_norm( reg_buffer_sz_bytes: int, epsilon: float, ) -> None: - torch.ops._C_custom_ar.sm70_tp2_all_reduce_gemma_rms_norm( + _custom_ar_op("sm70_tp2_all_reduce_gemma_rms_norm")( fa, inp, residual, @@ -3019,7 +3037,7 @@ def sm70_tp4_all_reduce_gemma_rms_norm( reg_buffer_sz_bytes: int, epsilon: float, ) -> None: - torch.ops._C_custom_ar.sm70_tp4_all_reduce_gemma_rms_norm( + _custom_ar_op("sm70_tp4_all_reduce_gemma_rms_norm")( fa, inp, residual, @@ -3044,7 +3062,7 @@ def sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( reg_buffer_sz_bytes: int, epsilon: float, ) -> None: - torch.ops._C_custom_ar.sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( + _custom_ar_op("sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather")( fa, inp, residual, @@ -3064,7 +3082,7 @@ def all_reduce_sum2( inp_b: torch.Tensor, out: torch.Tensor, ) -> None: - torch.ops._C_custom_ar.all_reduce_sum2(fa, inp_a, inp_b, out) + _custom_ar_op("all_reduce_sum2")(fa, inp_a, inp_b, out) def top1_argmax( @@ -3074,9 +3092,7 @@ def top1_argmax( reg_buffer: int, reg_buffer_sz_bytes: int, ) -> None: - torch.ops._C_custom_ar.top1_argmax( - fa, input_pair, out, reg_buffer, reg_buffer_sz_bytes - ) + _custom_ar_op("top1_argmax")(fa, input_pair, out, reg_buffer, reg_buffer_sz_bytes) def tile_runtime_all_reduce( @@ -3089,7 +3105,7 @@ def tile_runtime_all_reduce( engine_blocks: int, compute_iters: int, ) -> None: - torch.ops._C_custom_ar.tile_runtime_all_reduce( + _custom_ar_op("tile_runtime_all_reduce")( fa, inp, out, @@ -3112,7 +3128,7 @@ def tile_runtime_all_reduce_engine( reducer_blocks: int, compute_iters: int, ) -> None: - torch.ops._C_custom_ar.tile_runtime_all_reduce_engine( + _custom_ar_op("tile_runtime_all_reduce_engine")( fa, inp, out, @@ -3132,7 +3148,7 @@ def tile_runtime_wait_reduce( tile_numel: int, reducer_blocks: int, ) -> None: - torch.ops._C_custom_ar.tile_runtime_wait_reduce( + _custom_ar_op("tile_runtime_wait_reduce")( fa, staging, out, @@ -3142,54 +3158,54 @@ def tile_runtime_wait_reduce( def dispose(fa: int) -> None: - torch.ops._C_custom_ar.dispose(fa) + _custom_ar_op("dispose")(fa) def meta_size() -> int: - return torch.ops._C_custom_ar.meta_size() + return _custom_ar_op("meta_size")() def sm70_tp4_push_allreduce_buffer_size() -> int: - return torch.ops._C_custom_ar.sm70_tp4_push_allreduce_buffer_size() + return _custom_ar_op("sm70_tp4_push_allreduce_buffer_size")() def register_buffer(fa: int, ipc_tensors: list[int]) -> None: - return torch.ops._C_custom_ar.register_buffer(fa, ipc_tensors) + return _custom_ar_op("register_buffer")(fa, ipc_tensors) def register_sm70_tp4_push_allreduce_buffer(fa: int, ipc_tensors: list[int]) -> None: - torch.ops._C_custom_ar.register_sm70_tp4_push_allreduce_buffer(fa, ipc_tensors) + _custom_ar_op("register_sm70_tp4_push_allreduce_buffer")(fa, ipc_tensors) def get_graph_buffer_ipc_meta(fa: int) -> tuple[list[int], list[int]]: - return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa) + return _custom_ar_op("get_graph_buffer_ipc_meta")(fa) def register_graph_buffers( fa: int, handles: list[list[int]], offsets: list[list[int]] ) -> None: - torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets) + _custom_ar_op("register_graph_buffers")(fa, handles, offsets) def allocate_shared_buffer_and_handle(size: int) -> tuple[int, torch.Tensor]: - return torch.ops._C_custom_ar.allocate_shared_buffer_and_handle(size) + return _custom_ar_op("allocate_shared_buffer_and_handle")(size) def open_mem_handle(mem_handle: torch.Tensor): - return torch.ops._C_custom_ar.open_mem_handle(mem_handle) + return _custom_ar_op("open_mem_handle")(mem_handle) def free_shared_buffer(ptr: int) -> None: - torch.ops._C_custom_ar.free_shared_buffer(ptr) + _custom_ar_op("free_shared_buffer")(ptr) # quick all reduce def init_custom_qr(rank: int, world_size: int, qr_max_size: int | None = None) -> int: - return torch.ops._C_custom_ar.init_custom_qr(rank, world_size, qr_max_size) + return _custom_ar_op("init_custom_qr")(rank, world_size, qr_max_size) def qr_destroy(fa: int) -> None: - torch.ops._C_custom_ar.qr_destroy(fa) + _custom_ar_op("qr_destroy")(fa) def qr_all_reduce( @@ -3199,19 +3215,19 @@ def qr_all_reduce( quant_level: int, cast_bf2half: bool = False, ) -> None: - torch.ops._C_custom_ar.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half) + _custom_ar_op("qr_all_reduce")(fa, inp, out, quant_level, cast_bf2half) def qr_get_handle(fa: int) -> torch.Tensor: - return torch.ops._C_custom_ar.qr_get_handle(fa) + return _custom_ar_op("qr_get_handle")(fa) def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None: - return torch.ops._C_custom_ar.qr_open_handles(fa, handles) + return _custom_ar_op("qr_open_handles")(fa, handles) def qr_max_size() -> int: - return torch.ops._C_custom_ar.qr_max_size() + return _custom_ar_op("qr_max_size")() def get_flash_mla_metadata( diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index a554e2bf4d..4ebaecea8c 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -24,15 +24,28 @@ def _maybe_load_fp8_qpn8_library() -> None: return generic_override = os.getenv("VLLM_SM70_FP8_QPN8") specific_override = os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") + online_override = os.getenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8") generic_enabled = generic_override == "1" - default_route_enabled = generic_override != "0" and specific_override != "0" - if generic_enabled or default_route_enabled: + specific_enabled = generic_override != "0" and specific_override == "1" + online_enabled = online_override != "0" + if generic_enabled or specific_enabled or online_enabled: torch.ops.load_library(library_path) _maybe_load_fp8_qpn8_library() +def _maybe_load_nvfp4_qpn_m1_library() -> None: + """Load the narrow Qwen3.8 NVFP4 experiment in spawned TP workers.""" + library_path = os.getenv("VLLM_SM70_NVFP4_QPN_M1_LIBRARY") + route_enabled = os.getenv("VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE", "1") != "0" + if library_path is not None and route_enabled: + torch.ops.load_library(library_path) + + +_maybe_load_nvfp4_qpn_m1_library() + + def _maybe_load_sm70_sampler_library() -> None: """Load the sampler fragment only when the main ``vllm._C`` lacks it.""" if hasattr(torch.ops._C, "sm70_sample_chunked_top20_philox_token_out"): @@ -69,6 +82,26 @@ def _op(name: str): return getattr(torch.ops._C, name) +def _qwen38_qpn8_op(name: str): + """Prefer the task sidecar, then fall back to the production namespace.""" + sidecar = torch.ops._C_qwen38 + if hasattr(sidecar, name): + return getattr(sidecar, name) + return _op(name) + + +def has_fp8_qpn8_hc_dispatch() -> bool: + return hasattr(torch.ops._C_qwen38, "fp8_qpn8_hc_dispatch_sm70_out") or hasattr( + torch.ops._C, "fp8_qpn8_hc_dispatch_sm70_out" + ) + + +def has_nvfp4_qpn_m1_dispatch() -> bool: + return hasattr(torch.ops._C_qwen38, "nvfp4_moe_qpn_m1_sm70_out") or hasattr( + torch.ops._C, "nvfp4_moe_qpn_m1_sm70_out" + ) + + def silu_and_mul_interleaved(out: torch.Tensor, input: torch.Tensor) -> None: _op("silu_and_mul_interleaved")(out, input) @@ -490,12 +523,41 @@ def _fp8_gemm_sm70_prefill_prescaled_out_fake( return None +def fp8_gemm_sm70_prescaled_m1_out( + out: torch.Tensor, + input: torch.Tensor, + qweight: torch.Tensor, + prescaled_factors: torch.Tensor, + group_size: int, + k_ld: int, + q_ld: int, +) -> None: + _op("fp8_gemm_sm70_prescaled_m1_out")( + out, input, qweight, prescaled_factors, group_size, k_ld, q_ld + ) + + +if hasattr(torch.ops._C, "fp8_gemm_sm70_prescaled_m1_out"): + + @register_fake("_C::fp8_gemm_sm70_prescaled_m1_out") + def _fp8_gemm_sm70_prescaled_m1_out_fake( + out: torch.Tensor, + input: torch.Tensor, + qweight: torch.Tensor, + prescaled_factors: torch.Tensor, + group_size: int, + k_ld: int, + q_ld: int, + ) -> None: + return None + + def fp8_qpn8_prepare_sm70( qweight: torch.Tensor, scales: torch.Tensor, ) -> list[torch.Tensor]: """Pack checkpoint-native block/channel FP8 weights into QPN8 layout.""" - return _op("fp8_qpn8_prepare_sm70")(qweight, scales) + return _qwen38_qpn8_op("fp8_qpn8_prepare_sm70")(qweight, scales) if hasattr(torch.ops._C, "fp8_qpn8_prepare_sm70"): @@ -521,7 +583,7 @@ def fp8_qpn8_dequantize_sm70_out( group_scales: torch.Tensor, ) -> None: """Materialize one QPN8 weight into a caller-owned FP16 workspace.""" - _op("fp8_qpn8_dequantize_sm70_out")(out, codes, group_scales) + _qwen38_qpn8_op("fp8_qpn8_dequantize_sm70_out")(out, codes, group_scales) if hasattr(torch.ops._C, "fp8_qpn8_dequantize_sm70_out"): @@ -544,7 +606,7 @@ def fp8_qpn8_prefill_sm70_out( gated_silu: bool, ) -> None: """Dequantize QPN8 into bounded workspace and run a large-M FP16 GEMM.""" - _op("fp8_qpn8_prefill_sm70_out")( + _qwen38_qpn8_op("fp8_qpn8_prefill_sm70_out")( out, dense_weight_ptr, input, @@ -580,7 +642,7 @@ def fp8_qpn8_dispatch_sm70_out( gated_silu: bool, ) -> None: """Runtime-dispatch dynamic M without specializing a Python branch.""" - _op("fp8_qpn8_dispatch_sm70_out")( + _qwen38_qpn8_op("fp8_qpn8_dispatch_sm70_out")( out, dense_weight_ptr, input, @@ -625,7 +687,7 @@ def fp8_qpn8_gemm_sm70_out( The model- and shape-gated automatic route and operator benchmark share this entry point. ``codes`` and ``group_scales`` use the QPN8 layout. """ - _op("fp8_qpn8_gemm_sm70_out")( + _qwen38_qpn8_op("fp8_qpn8_gemm_sm70_out")( out, input, codes, @@ -664,7 +726,7 @@ def fp8_qpn8_gemm_ba_split_sm70_out( ba_weight: torch.Tensor, ) -> None: """Run the exact-shape GDN QKV/Z FP8 and b/a FP16 projections.""" - _op("fp8_qpn8_gemm_ba_split_sm70_out")( + _qwen38_qpn8_op("fp8_qpn8_gemm_ba_split_sm70_out")( qkv_out, z_out, b_out, @@ -706,7 +768,7 @@ def fp8_qpn8_dispatch_ba_split_sm70_out( ba_weight: torch.Tensor, ) -> None: """Dispatch exact M=1 fusion or the original large-M projection path.""" - _op("fp8_qpn8_dispatch_ba_split_sm70_out")( + _qwen38_qpn8_op("fp8_qpn8_dispatch_ba_split_sm70_out")( qkv_out, z_out, b_out, @@ -751,7 +813,7 @@ def fp8_qpn8_gated_pair_sm70_out( prefetch_codes: bool = False, ) -> None: """Run the single-kernel paired-tile QPN8 gated SiLU experiment.""" - _op("fp8_qpn8_gated_pair_sm70_out")( + _qwen38_qpn8_op("fp8_qpn8_gated_pair_sm70_out")( out, input, codes, @@ -763,6 +825,77 @@ def fp8_qpn8_gated_pair_sm70_out( ) +def fp8_qpn8_hc_dispatch_sm70_out( + block_out: torch.Tensor, + injection_out: torch.Tensor, + down_staging: torch.Tensor, + lora_staging: torch.Tensor, + gate_staging: torch.Tensor, + partials: torch.Tensor, + dense_weight_ptr: int, + xn: torch.Tensor, + down_codes: torch.Tensor, + down_scales: torch.Tensor, + up_codes: torch.Tensor, + up_scales: torch.Tensor, +) -> None: + """Dispatch the exact Qwen4Exp HC pair without a Python M branch.""" + _qwen38_qpn8_op("fp8_qpn8_hc_dispatch_sm70_out")( + block_out, + injection_out, + down_staging, + lora_staging, + gate_staging, + partials, + dense_weight_ptr, + xn, + down_codes, + down_scales, + up_codes, + up_scales, + ) + + +if hasattr(torch.ops._C, "fp8_qpn8_hc_dispatch_sm70_out"): + + @register_fake("_C::fp8_qpn8_hc_dispatch_sm70_out") + def _fp8_qpn8_hc_dispatch_sm70_out_fake( + block_out: torch.Tensor, + injection_out: torch.Tensor, + down_staging: torch.Tensor, + lora_staging: torch.Tensor, + gate_staging: torch.Tensor, + partials: torch.Tensor, + dense_weight_ptr: int, + xn: torch.Tensor, + down_codes: torch.Tensor, + down_scales: torch.Tensor, + up_codes: torch.Tensor, + up_scales: torch.Tensor, + ) -> None: + return None + + +if hasattr(torch.ops._C_qwen38, "fp8_qpn8_hc_dispatch_sm70_out"): + + @register_fake("_C_qwen38::fp8_qpn8_hc_dispatch_sm70_out") + def _fp8_qpn8_hc_dispatch_sm70_out_sidecar_fake( + block_out: torch.Tensor, + injection_out: torch.Tensor, + down_staging: torch.Tensor, + lora_staging: torch.Tensor, + gate_staging: torch.Tensor, + partials: torch.Tensor, + dense_weight_ptr: int, + xn: torch.Tensor, + down_codes: torch.Tensor, + down_scales: torch.Tensor, + up_codes: torch.Tensor, + up_scales: torch.Tensor, + ) -> None: + return None + + if hasattr(torch.ops._C, "fp8_qpn8_gated_pair_sm70_out"): @register_fake("_C::fp8_qpn8_gated_pair_sm70_out") @@ -779,6 +912,29 @@ def _fp8_qpn8_gated_pair_sm70_out_fake( return None +def _register_qwen38_qpn8_out_fakes() -> None: + """Teach Dynamo about out-variant operators from the task sidecar.""" + + def fake_out(*args, **kwargs) -> None: + del args, kwargs + return None + + for name in ( + "fp8_qpn8_dequantize_sm70_out", + "fp8_qpn8_prefill_sm70_out", + "fp8_qpn8_dispatch_sm70_out", + "fp8_qpn8_gemm_sm70_out", + "fp8_qpn8_gemm_ba_split_sm70_out", + "fp8_qpn8_dispatch_ba_split_sm70_out", + "fp8_qpn8_gated_pair_sm70_out", + ): + if hasattr(torch.ops._C_qwen38, name): + register_fake(f"_C_qwen38::{name}")(fake_out) + + +_register_qwen38_qpn8_out_fakes() + + def nvfp4_qpn4_prepare_sm70( qweight: torch.Tensor, scales: torch.Tensor, @@ -1193,6 +1349,83 @@ def _mxfp4_moe_dense_stage_sm70_out_fake( return None +def mxfp4_moe_qpn_m1_sm70_out( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, + broadcast_input: bool, +) -> None: + _op("mxfp4_moe_qpn_m1_sm70_out")( + out, input, weights, scales, expert_ids, broadcast_input + ) + + +if hasattr(torch.ops._C, "mxfp4_moe_qpn_m1_sm70_out"): + + @register_fake("_C::mxfp4_moe_qpn_m1_sm70_out") + def _mxfp4_moe_qpn_m1_sm70_out_fake( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, + broadcast_input: bool, + ) -> None: + return None + + +def nvfp4_moe_qpn_m1_sm70_out( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, + broadcast_input: bool, + split_k: int, +) -> None: + _qwen38_qpn8_op("nvfp4_moe_qpn_m1_sm70_out")( + out, + input, + weights, + scales, + expert_ids, + broadcast_input, + split_k, + ) + + +if hasattr(torch.ops._C, "nvfp4_moe_qpn_m1_sm70_out"): + + @register_fake("_C::nvfp4_moe_qpn_m1_sm70_out") + def _nvfp4_moe_qpn_m1_sm70_out_fake( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, + broadcast_input: bool, + split_k: int, + ) -> None: + return None + + +if hasattr(torch.ops._C_qwen38, "nvfp4_moe_qpn_m1_sm70_out"): + + @register_fake("_C_qwen38::nvfp4_moe_qpn_m1_sm70_out") + def _nvfp4_moe_qpn_m1_sm70_out_sidecar_fake( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, + broadcast_input: bool, + split_k: int, + ) -> None: + return None + + def nvfp4_moe_dense_stage_sm70_out( out: torch.Tensor, input: torch.Tensor, @@ -1463,6 +1696,65 @@ def _sm70_f16_gemm_out_fake( return None +def sm70_glm_mhc_pre_norm_out( + gemm_mul: torch.Tensor, + gemm_sqrsum: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + layer_input: torch.Tensor, + norm_weight: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult: float, + sinkhorn_repeat: int, + norm_eps: float, +) -> None: + _op("sm70_glm_mhc_pre_norm_out")( + gemm_mul, + gemm_sqrsum, + hc_scale, + hc_base, + residual, + post_mix, + comb_mix, + layer_input, + norm_weight, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult, + sinkhorn_repeat, + norm_eps, + ) + + +if hasattr(torch.ops._C, "sm70_glm_mhc_pre_norm_out"): + + @register_fake("_C::sm70_glm_mhc_pre_norm_out") + def _sm70_glm_mhc_pre_norm_out_fake( + gemm_mul: torch.Tensor, + gemm_sqrsum: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + layer_input: torch.Tensor, + norm_weight: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult: float, + sinkhorn_repeat: int, + norm_eps: float, + ) -> None: + return None + + def sm70_f16_indexed_rerank_out( out: torch.Tensor, input: torch.Tensor, @@ -1510,6 +1802,31 @@ def _sm70_f16_indexed_rerank_out_fake( return None +def sm70_glm_kda_fg_b_out( + f_out: torch.Tensor, + g_out: torch.Tensor, + f_input: torch.Tensor, + g_input: torch.Tensor, + f_weight: torch.Tensor, + g_weight: torch.Tensor, +) -> None: + _op("sm70_glm_kda_fg_b_out")(f_out, g_out, f_input, g_input, f_weight, g_weight) + + +if hasattr(torch.ops._C, "sm70_glm_kda_fg_b_out"): + + @register_fake("_C::sm70_glm_kda_fg_b_out") + def _sm70_glm_kda_fg_b_out_fake( + f_out: torch.Tensor, + g_out: torch.Tensor, + f_input: torch.Tensor, + g_input: torch.Tensor, + f_weight: torch.Tensor, + g_weight: torch.Tensor, + ) -> None: + return None + + def sm70_f16_indexed_rerank_packed_out( out: torch.Tensor, input: torch.Tensor, diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index 08dfbd4937..97ce957f00 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -44,8 +44,7 @@ def get_op_overload( ROTARY_OP = get_op_overload(torch.ops._C, "rotary_embedding") -FLASHINFER_ROTARY_OP = get_op_overload(torch.ops.vllm, - "flashinfer_rotary_embedding") +FLASHINFER_ROTARY_OP = get_op_overload(torch.ops.vllm, "flashinfer_rotary_embedding") QUANT_OPS: dict[QuantKey, OpOverload] = {} diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py index 01f5d5a210..9fef59b00b 100644 --- a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py +++ b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py @@ -198,8 +198,7 @@ def __init__(self, config: VllmConfig) -> None: dtype = config.model_config.dtype if FUSED_QK_ROPE_OP is None: logger.warning_once( - "QK Norm+RoPE fusion not enabled: fused_qk_norm_rope is " - "unavailable." + "QK Norm+RoPE fusion not enabled: fused_qk_norm_rope is unavailable." ) return diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py index b2c465b1ff..d5de88ba2c 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py @@ -124,12 +124,9 @@ def _register_fused_op( _register_fused_op(kFp8StaticTensorSym, False, "rms_norm_static_fp8_quant") -_register_fused_op(kFp8StaticTensorSym, True, - "fused_add_rms_norm_static_fp8_quant") -_register_fused_op(kFp8DynamicTokenSym, False, - "rms_norm_dynamic_per_token_quant") -_register_fused_op(kFp8DynamicTokenSym, True, - "rms_norm_dynamic_per_token_quant") +_register_fused_op(kFp8StaticTensorSym, True, "fused_add_rms_norm_static_fp8_quant") +_register_fused_op(kFp8DynamicTokenSym, False, "rms_norm_dynamic_per_token_quant") +_register_fused_op(kFp8DynamicTokenSym, True, "rms_norm_dynamic_per_token_quant") _register_fused_op(kFp8Dynamic128Sym, False, "rms_norm_per_block_quant") _register_fused_op(kFp8Dynamic128Sym, True, "rms_norm_per_block_quant") _register_fused_op(kFp8Dynamic64Sym, False, "rms_norm_per_block_quant") diff --git a/vllm/compilation/passes/utility/fix_functionalization.py b/vllm/compilation/passes/utility/fix_functionalization.py index 3f955d53cf..3e8c3af1c0 100644 --- a/vllm/compilation/passes/utility/fix_functionalization.py +++ b/vllm/compilation/passes/utility/fix_functionalization.py @@ -78,9 +78,7 @@ def __call__(self, graph: torch.fx.Graph) -> None: fused_add_rms_norm_static_fp8_quant = _get_c_op( "fused_add_rms_norm_static_fp8_quant" ) - rms_norm_dynamic_per_token_quant = _get_c_op( - "rms_norm_dynamic_per_token_quant" - ) + rms_norm_dynamic_per_token_quant = _get_c_op("rms_norm_dynamic_per_token_quant") rms_norm_targets = [ op for op in ( diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 2b0833e545..728fad8e8e 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -759,6 +759,8 @@ class CompilationConfig: "vllm::mamba_mixer2", "vllm::mamba_mixer", "vllm::short_conv", + "vllm::qwen4_exp_ple_short_conv", + "vllm::qwen4_exp_qsa_with_output", "vllm::linear_attention", "vllm::plamo2_mamba_mixer", "vllm::qwen_gdn_attention_core", @@ -775,6 +777,7 @@ class CompilationConfig: "vllm::olmo_hybrid_gdn_full_forward", "vllm::kda_attention", "vllm::sparse_attn_indexer", + "vllm::sparse_attn_indexer_kpool", "vllm::rocm_aiter_sparse_attn_indexer", "vllm::deepseek_v4_attention", ] diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index b89d8b0210..6cfd519856 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -45,6 +45,7 @@ "exaone_moe_mtp", "exaone4_5_mtp", "qwen3_next_mtp", + "qwen4_exp_mtp", "qwen3_5_mtp", "longcat_flash_mtp", "mtp", @@ -81,6 +82,38 @@ DraftSampleMethod = Literal["greedy", "probabilistic"] +def get_dflash_model_draft_tokens(speculative_config: Any) -> int: + """Return the block width produced by the DFlash checkpoint itself. + + Lookup-augmented DFlash2 may hand a wider proposal to the target, but its + model/convolutions must retain the block size used during training. + """ + verify_tokens = int(getattr(speculative_config, "num_speculative_tokens", 0) or 0) + if getattr(speculative_config, "method", None) != "dflash" or not getattr( + speculative_config, "ngram_assist", False + ): + return verify_tokens + + draft_model_config = getattr(speculative_config, "draft_model_config", None) + hf_config = getattr(draft_model_config, "hf_config", None) + dflash_config = getattr(hf_config, "dflash_config", None) or {} + trained_tokens = int(dflash_config.get("block_size", 0) or 0) - 1 + has_selector = int(dflash_config.get("selector_top_k", 0) or 0) > 0 + if has_selector and 0 < trained_tokens < verify_tokens: + return trained_tokens + return verify_tokens + + +def uses_adaptive_dflash_lookup(speculative_config: Any) -> bool: + """Whether DFlash2 can switch between model and augmented widths.""" + verify_tokens = int(getattr(speculative_config, "num_speculative_tokens", 0) or 0) + return ( + getattr(speculative_config, "method", None) == "dflash" + and bool(getattr(speculative_config, "ngram_assist", False)) + and 0 < get_dflash_model_draft_tokens(speculative_config) < verify_tokens + ) + + @config class SpeculativeConfig: """Configuration for speculative decoding.""" @@ -138,6 +171,9 @@ class SpeculativeConfig: """The specific revision to use for the draft model code on Hugging Face Hub. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version.""" + index_share_for_mtp_iteration: bool | None = None + """Override whether MTP iterations reuse the first step's sparse indices. + If ``None``, use the value from the draft model's Hugging Face config.""" # Advanced control disable_padded_drafter_batch: bool = False @@ -159,10 +195,12 @@ class SpeculativeConfig: """Minimum size of ngram token window when using Ngram proposer, if provided. Defaults to 1.""" ngram_assist: bool = False - """Try prompt-ngram lookup before DFlash2 draft generation. Full-width - ngram hits skip the DFlash2 query and selector while preserving its - context-KV state. Only valid with ``method='dflash'`` and a DFlash2 - selector capability.""" + """Combine prompt-ngram lookup with DFlash2. At the checkpoint-native + width, full hits may skip the DFlash2 query and selector. If + ``num_speculative_tokens`` is wider than the checkpoint block, DFlash2 + keeps its trained width and lookup may fill the extra target-verification + positions. Only valid with ``method='dflash'`` and a DFlash2 selector + capability.""" # Alternative drafting strategies parallel_drafting: bool = False @@ -367,6 +405,15 @@ def compute_hash(self) -> str: ) ) + if self.method == "mtp" and self.draft_model_config is not None: + factors.append( + getattr( + self.draft_model_config.hf_config, + "index_share_for_mtp_iteration", + False, + ) + ) + hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() return hash_str @@ -516,6 +563,26 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: {"n_predict": n_predict, "architectures": ["Qwen3NextMTP"]} ) + if hf_config.model_type in {"qwen4_exp", "qwen4_exp_text"}: + hf_config.model_type = "qwen4_exp_mtp" + if hf_config.model_type == "qwen4_exp_mtp": + text_config = get_hf_text_config(hf_config) + n_predict = getattr( + text_config, + "mtp_num_hidden_layers", + getattr(text_config, "num_nextn_predict_layers", None), + ) + hf_config.update( + { + "hc_mult": int(text_config.hc_count), + "n_predict": n_predict, + "architectures": ["Qwen4ExpMTP"], + "index_share_for_mtp_iteration": getattr( + text_config, "index_share_for_mtp_iteration", True + ), + } + ) + if hf_config.model_type == "exaone_moe": hf_config.model_type = "exaone_moe_mtp" if hf_config.model_type == "exaone_moe_mtp": @@ -968,6 +1035,15 @@ def __post_init__(self): ), ) ) + + if self.index_share_for_mtp_iteration is not None: + if self.method != "mtp" or self.draft_model_config is None: + raise ValueError( + "index_share_for_mtp_iteration is only supported with method='mtp'" + ) + self.draft_model_config.hf_config.index_share_for_mtp_iteration = ( + self.index_share_for_mtp_iteration + ) return self def _verify_dspark_final_stage_ownership(self) -> None: @@ -999,7 +1075,7 @@ def _validate_suffix_decoding(self): if not has_arctic_inference(): raise ImportError( "Arctic Inference is required for suffix decoding. " - "Install via `pip install arctic-inference==0.1.1`." + "Install via `pip install arctic-inference==0.2.0`." ) if self.num_speculative_tokens is None: # Suffix decoding decides the actual number of speculative tokens @@ -1284,6 +1360,15 @@ def use_step3p5_mtp(self) -> bool: == "step3p5_mtp" ) + def use_qwen4_exp_mtp(self) -> bool: + """Return whether Qwen4Exp needs its dedicated proposer.""" + return ( + self.method == "mtp" + and self.draft_model_config is not None + and getattr(self.draft_model_config.hf_config, "model_type", None) + == "qwen4_exp_mtp" + ) + def use_eagle(self) -> bool: return ( self.method in ("eagle", "eagle3", "mtp") diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 961b979c12..6e918743e9 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -47,7 +47,12 @@ from .profiler import ProfilerConfig from .reasoning import ReasoningConfig from .scheduler import SchedulerConfig -from .speculative import EagleModelTypes, NgramGPUTypes, SpeculativeConfig +from .speculative import ( + EagleModelTypes, + NgramGPUTypes, + SpeculativeConfig, + uses_adaptive_dflash_lookup, +) from .structured_outputs import StructuredOutputsConfig from .utils import SupportsHash, config, replace from .weight_transfer import WeightTransferConfig @@ -66,7 +71,15 @@ logger = init_logger(__name__) -DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset({"Qwen3ForCausalLM"}) +DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( + { + "Glm5NextForCausalLM", + "Glm5NextForConditionalGeneration", + "Qwen3ForCausalLM", + "Qwen4ExpForCausalLM", + "Qwen4ExpForConditionalGeneration", + } +) _SM70_NOMTP_CUDAGRAPH_CAPTURE_SIZES = (1, 2, 4, 8, 16) _SM70_MTP_CUDAGRAPH_REQUEST_SIZES = (1, 2, 4, 6, 8, 12, 16) @@ -693,11 +706,18 @@ def _is_default_v2_model_runner_model(self) -> bool: return False architectures = getattr(model_config, "architectures", []) - if not any( + is_default_v2_architecture = any( arch in DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES for arch in architectures - ): + ) + if not is_default_v2_architecture: return False + # Qwen4Exp is a hybrid MoE model whose QSA ring cache and PLE raw-token + # inputs are implemented by Model Runner V2. Its quantized checkpoint + # is therefore an explicit V2 route rather than a generic fallback. + if any(arch.startswith("Qwen4ExpFor") for arch in architectures): + return True + return not model_config.is_moe and not model_config.is_quantized @property @@ -1110,11 +1130,23 @@ def __post_init__(self): and self.speculative_config.use_dflash_ddtree() and not self.speculative_config.ddtree_disable_tree_verify ) + adaptive_dflash_lookup = bool( + self.speculative_config is not None + and uses_adaptive_dflash_lookup(self.speculative_config) + and envs.VLLM_DFLASH2_LOOKUP_ADAPTIVE + ) if self.scheduler_config.async_scheduling: # Async scheduling explicitly enabled, hard fail any incompatibilities. # Currently, async scheduling only support eagle speculative # decoding. + if adaptive_dflash_lookup: + raise ValueError( + "Async scheduling fixes the DFlash2 proposal width at its " + "configured maximum and is incompatible with adaptive q8/q16 " + "lookup verification. Disable async scheduling or set " + "VLLM_DFLASH2_LOOKUP_ADAPTIVE=0 to use fixed q16 verification." + ) if dflash_ddtree_tree_verify: raise ValueError( "Async scheduling is not compatible with dflash_ddtree " @@ -1153,6 +1185,13 @@ def __post_init__(self): "Disabling asynchronous scheduling by default for pooling model." ) self.scheduler_config.async_scheduling = False + elif adaptive_dflash_lookup: + logger.warning_once( + "Disabling asynchronous scheduling because adaptive DFlash2 " + "lookup verification must return a per-step q8/q16 width to " + "the scheduler." + ) + self.scheduler_config.async_scheduling = False elif dflash_ddtree_tree_verify: logger.warning_once( "Async scheduling is not compatible with dflash_ddtree tree " diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 8d55ceeaff..b0bb55b7cb 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -339,7 +339,8 @@ def __init__( ) logger.info( "SM70 TP4 SGLang-style push all-reduce enabled for the " - "FP16 80-KiB verifier and 8-KiB decode payloads." + "FP16 80-KiB verifier, 8-KiB decode, and 5-KiB Qwen4Exp " + "payloads." ) @contextmanager @@ -380,6 +381,8 @@ def register_graph_buffers(self): def should_custom_ar(self, inp: torch.Tensor): if self.disabled: return False + if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16): + return False inp_size = inp.numel() * inp.element_size() # custom allreduce requires input byte size to be multiples of 16 if inp_size % 16 != 0: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 31850497d3..4dc0954a72 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1819,10 +1819,9 @@ def _maybe_apply_sm70_mtp_defaults( return if "draft_sample_method" not in self.speculative_config: - # Preserve the validated probabilistic default for MTP and the - # DFlash family. DFlash2's greedy performance lane remains - # available when explicitly requested, without changing DFlash1. - draft_sample_method = "probabilistic" + # MTP can use the vocab-parallel local-argmax fast path only in + # greedy mode. DFlash keeps probabilistic drafting by default. + draft_sample_method = "greedy" if spec_method == "mtp" else "probabilistic" self.speculative_config["draft_sample_method"] = draft_sample_method profile_updates.append( f"speculative_config.draft_sample_method={draft_sample_method}" diff --git a/vllm/envs.py b/vllm/envs.py index 443e004a1a..9940e5b3ad 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -159,19 +159,29 @@ VLLM_SM70_FP8_GROUPED_BMM_DECODE: bool = True VLLM_SM70_FP8_PREFILL_FAST_SELECTOR: bool = True VLLM_SM70_FP8_PREFILL_PRESCALED: bool = True + VLLM_SM70_FP8_PRESCALED_M1_DECODE: bool = True VLLM_SM70_FP8_PREFILL_CUTLASS: bool = True VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS: bool = True VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS_ONLY: bool = False VLLM_SM70_FP8_PREFILL_EXACT_DENSE: bool = True VLLM_SM70_FP8_QPN8: bool = False - VLLM_SM70_FP8_QPN8_PP2_TP4: bool = True + VLLM_SM70_QWEN4_EXP_ONLINE_QPN8: bool = True + VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION: bool = True + VLLM_SM70_FP8_QPN8_PP2_TP4: bool = False + VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE: bool = False VLLM_SM70_FP8_QPN8_LIBRARY: str | None = None VLLM_SM70_SAMPLER_LIBRARY: str | None = None + VLLM_SM70_FA2_D256_LIBRARY: str | None = None VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM: bool = False VLLM_SM70_NVFP4_QPN2: bool = False + VLLM_SM70_NVFP4_QPN2_PREFILL: bool = False + VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M: int = 1024 VLLM_SM70_MXFP4_TUNE_SMALL_SHAPES: bool = True VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES: bool = True VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR: bool = True + VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE: bool = True + VLLM_SM70_NVFP4_QPN_M1_LIBRARY: str | None = None + VLLM_SM70_QWEN38_ROUTER_TOPK: bool = True VLLM_SM70_AWQ_REUSE_IMPORTED_CACHE: bool = False VLLM_SM70_AWQ_WARMUP: bool = True VLLM_SM70_AWQ_WARMUP_MAX_M: int = 16 @@ -183,6 +193,7 @@ VLLM_SM70_MXFP4_DENSE_TUNE_MAX_M: int = 16 VLLM_SM70_NVFP4_DENSE_TUNE_MAX_M: int = 16 VLLM_SM70_DSV4_FP16_GEMV: bool = False + VLLM_SM70_DSV4_FP13_GEMV: bool = True VLLM_SM70_DSV4_MHC_FP32_STAGE: bool = True VLLM_SM70_DSV4_QNORM_KV_FUSED_TP4: bool = True VLLM_SM70_PP_STATIC_HIDDEN_TRANSFER: bool = True @@ -195,6 +206,7 @@ VLLM_SM70_DFLASH2_QPN8_RERANK: bool = False VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW: bool = False VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER: bool = True + VLLM_SM70_DFLASH2_QPN8_ALLOW_CANDIDATE_ORDER: bool = False VLLM_SM70_DFLASH2_VERIFY_FASTPATH: bool = False VLLM_SM70_DFLASH2_FUSED_GDN_METADATA: bool = False VLLM_SM70_DFLASH2_GDN_METADATA_SHADOW: bool = False @@ -209,6 +221,7 @@ VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION: bool = False VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC: bool = False VLLM_SM70_TP4_PUSH_ALLREDUCE: bool = True + VLLM_SM70_CUSTOM_AR_LIBRARY: str | None = None VLLM_SM70_TOP1_CUSTOM_AR: bool = False VLLM_SM70_GREEDY_TOKEN_FASTPATH: bool = True VLLM_SM70_GREEDY_TOKEN_FASTPATH_TRACE: bool = False @@ -246,13 +259,14 @@ VLLM_SM70_MXFP4_TURBOMIND: bool = True VLLM_SM70_MXFP4_MOE_ACTIVE_EXPERT_B1: bool = False VLLM_SM70_MXFP4_MOE_ACTIVE_EXPERT_MAX_TOKENS: int = 8 - VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE: bool = False + VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE: bool = True VLLM_SM70_MXFP4_MOE_GROUPED_M8: bool = False VLLM_SM70_MXFP4_MOE_GROUPED_VERIFIER: bool = False VLLM_SM70_MXFP4_MOE_GROUPED_M8_EXPERT_ROWS: bool = False VLLM_SM70_MXFP4_MOE_GROUPED_M8_FAST_SELECTOR: bool = True - VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE: bool = False - VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE: bool = False + VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE: bool = True + VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE: bool = True + VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE: bool = True VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE: bool = True VLLM_SM70_DSV4_SPARSE_MLA_SPLITK_SWA: bool = False VLLM_SM70_DSV4_SPARSE_MLA_SPLITK_C4: bool = False @@ -390,8 +404,10 @@ VLLM_FLASH_V100_DECODE_USE_BHMD_OUT: bool = True VLLM_FLASH_V100_DECODE_USE_WMMA_WRAPPER: bool = False VLLM_FLASH_V100_DECODE_USE_XQA: bool = True - VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY: bool = False + VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY: bool = True VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY_MIN_MODEL_LEN: int = 32768 + VLLM_FLASH_V100_DFLASH2_FIXED_INTERLEAVED: bool = True + VLLM_FLASH_V100_DFLASH2_STAGE_PAGE_IDS: bool = True VLLM_FLASH_V100_DECODE_XQA_Q4_MIN_SEQ_LEN: int = 32768 VLLM_FLASH_V100_XQA_G6_P1024_SAWTOOTH: bool = True VLLM_FLASH_V100_XQA_G6_P1024_SAWTOOTH_TRACE: bool = False @@ -1686,17 +1702,33 @@ def _resolve_rust_frontend_path() -> str | None: # a mixed NVFP4 checkpoint may select its separately validated default in # the compressed-tensors scheme. Explicit 0 disables both routes. "VLLM_SM70_FP8_QPN8": lambda: bool(int(os.getenv("VLLM_SM70_FP8_QPN8", "0"))), - # Default-on QPN8 route for the validated serialized PP2 x TP4 contract. - # Admission additionally requires exact operator shapes/layouts, B1, - # no speculative decoding, no DBO, and no explicit ubatching. Set this to - # 0 (or the generic QPN8 flag to 0) to retain TurboMind everywhere. + "VLLM_SM70_QWEN4_EXP_ONLINE_QPN8": lambda: bool( + int(os.getenv("VLLM_SM70_QWEN4_EXP_ONLINE_QPN8", "1")) + ), + # Exact M=1 Qwen3Next/Qwen4Exp shared-expert output gate. This replaces + # the scalar GEMV, sigmoid, and output multiply with one SM70 kernel while + # retaining the checkpoint's FP16 accumulation and output rounding. + "VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION": lambda: bool( + int(os.getenv("VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION", "1")) + ), + # Experimental QPN8 route for the serialized PP2 x TP4 contract. It is + # default-off after matched model-level quality regressions. An explicit + # opt-in still requires exact operator shapes/layouts, B1, no speculative + # decoding, no DBO, and no explicit ubatching. "VLLM_SM70_FP8_QPN8_PP2_TP4": lambda: bool( - int(os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "1")) + int(os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4", "0")) + ), + # Experimental non-fused QPN8 route for the exact PP2 x TP4 shared-expert + # gate/up tensor. The model-level clamp-SwiGLU remains external. This + # numerically sensitive role requires an explicit opt-in. + "VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE", "0")) ), # Optional source-built QPN8-only extension. Production builds leave this # unset because the same operators are linked into vllm._C. "VLLM_SM70_FP8_QPN8_LIBRARY": lambda: os.getenv("VLLM_SM70_FP8_QPN8_LIBRARY", None), "VLLM_SM70_SAMPLER_LIBRARY": lambda: os.getenv("VLLM_SM70_SAMPLER_LIBRARY", None), + "VLLM_SM70_FA2_D256_LIBRARY": lambda: os.getenv("VLLM_SM70_FA2_D256_LIBRARY", None), "VLLM_SM70_FP8_PREFILL_CUTLASS": lambda: bool( int(os.getenv("VLLM_SM70_FP8_PREFILL_CUTLASS", "1")) ), @@ -1706,6 +1738,15 @@ def _resolve_rust_frontend_path() -> str | None: # QPN2 is an explicit opt-in for compatible NVFP4 small-M shapes; larger M # stays on the existing TurboMind path. "VLLM_SM70_NVFP4_QPN2": lambda: bool(int(os.getenv("VLLM_SM70_NVFP4_QPN2", "0"))), + # Reuse the already resident QPN2 code/scale layout for bounded-workspace + # FP16 large-M prefill. M<=8 decode and speculative verification remain on + # QPN2. This stays opt-in until full-model speed and quality gates pass. + "VLLM_SM70_NVFP4_QPN2_PREFILL": lambda: bool( + int(os.getenv("VLLM_SM70_NVFP4_QPN2_PREFILL", "0")) + ), + "VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M": lambda: int( + os.getenv("VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M", "1024") + ), # Experimental TileRT-inspired down-proj lane: after the row-parallel AWQ # GEMM, use the local tile-runtime TP2 all-reduce substrate for the MLP # hidden-state reduction. This is default-off until it wins end-to-end. @@ -1766,6 +1807,9 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_FP8_PREFILL_PRESCALED": lambda: bool( int(os.getenv("VLLM_SM70_FP8_PREFILL_PRESCALED", "1")) ), + "VLLM_SM70_FP8_PRESCALED_M1_DECODE": lambda: bool( + int(os.getenv("VLLM_SM70_FP8_PRESCALED_M1_DECODE", "1")) + ), "VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS": lambda: bool( int(os.getenv("VLLM_SM70_FP8_PRESERVE_DEFAULT_SPLITS", "1")) ), @@ -1781,6 +1825,21 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR": lambda: bool( int(os.getenv("VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR", "1")) ), + # Direct ten-route M=1 expert GEMMs for Qwen3.8 Flash Next NVFP4 TP4. + # This consumes the existing TurboMind-packed weights while skipping the + # replicated-input prepare kernel. The fixed TP4 8192x512 acceptance run + # reached 82.274 tokens/s with unchanged arithmetic and Chinese hashes. + "VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE": lambda: bool( + int(os.getenv("VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE", "1")) + ), + "VLLM_SM70_NVFP4_QPN_M1_LIBRARY": lambda: os.getenv( + "VLLM_SM70_NVFP4_QPN_M1_LIBRARY" + ), + # Exact single-token E=512, K=10 softmax router used by Qwen3.8 Flash + # Next. All other shapes and scoring modes retain the generic CUDA op. + "VLLM_SM70_QWEN38_ROUTER_TOPK": lambda: bool( + int(os.getenv("VLLM_SM70_QWEN38_ROUTER_TOPK", "1")) + ), "VLLM_SM70_AWQ_REUSE_IMPORTED_CACHE": lambda: bool( int(os.getenv("VLLM_SM70_AWQ_REUSE_IMPORTED_CACHE", "0")) ), @@ -1813,6 +1872,9 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_DSV4_FP16_GEMV": lambda: bool( int(os.getenv("VLLM_SM70_DSV4_FP16_GEMV", "0")) ), + "VLLM_SM70_DSV4_FP13_GEMV": lambda: bool( + int(os.getenv("VLLM_SM70_DSV4_FP13_GEMV", "1")) + ), "VLLM_SM70_DSV4_MHC_FP32_STAGE": lambda: bool( int(os.getenv("VLLM_SM70_DSV4_MHC_FP32_STAGE", "1")) ), @@ -1884,6 +1946,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER": lambda: bool( int(os.getenv("VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER", "1")) ), + # Candidate-order tie handling is a benchmark-only experiment. Requiring + # a second opt-in prevents stale deployment scripts from silently trading + # scored quality for a small selector win. + "VLLM_SM70_DFLASH2_QPN8_ALLOW_CANDIDATE_ORDER": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_QPN8_ALLOW_CANDIDATE_ORDER", "0")) + ), # Umbrella gate for selector-based DFlash verification optimizations. Keep # this default-off while each stage is checked against the unchanged path. "VLLM_SM70_DFLASH2_VERIFY_FASTPATH": lambda: bool( @@ -1969,6 +2037,9 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_TP4_PUSH_ALLREDUCE": lambda: bool( int(os.getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE", "1")) ), + # Optional task-built custom-AR fragment. Operators present in the sidecar + # override the production namespace; every other operator falls back. + "VLLM_SM70_CUSTOM_AR_LIBRARY": lambda: os.getenv("VLLM_SM70_CUSTOM_AR_LIBRARY"), # Safe greedy-only shortcut: avoid full vocab all-gather/sampler work when # the request batch is pure greedy and has no penalties, logprobs, grammar, # or custom logits processing. It still computes local logits with the @@ -2177,7 +2248,7 @@ def _resolve_rust_frontend_path() -> str | None: # Fuse the six one-row DeepSeek V4 MXFP4 decode experts into one # TurboMind launch. The C++ route reads this value directly as well. "VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE": lambda: bool( - int(os.getenv("VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE", "0")) + int(os.getenv("VLLM_SM70_MXFP4_MOE_COMPACT_GROUPED_DECODE", "1")) ), # Experimental verifier-M8 grouped dispatch. This collapses the fixed 48 # active-expert stage calls into one TurboMind grouped launch. @@ -2203,14 +2274,19 @@ def _resolve_rust_frontend_path() -> str | None: # Skip the generic 256-expert sort/permute/unpermute pipeline for the # exact DeepSeek V4 B1, replicated-expert, top-k=6 decode contract. "VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE": lambda: bool( - int(os.getenv("VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE", "0")) + int(os.getenv("VLLM_SM70_MXFP4_MOE_DIRECT_TOP6_DECODE", "1")) ), # Keep the six B1 routes in their original top-k order. Compact W13/W2 # then consume topk_ids directly, so no sort/inverse-permutation prepare - # kernel is needed. This stays independently gated until endpoint and - # quality acceptance are complete. + # kernel is needed. Exact shape/route checks remain in the caller and =0 + # restores the stable-sort path. "VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE": lambda: bool( - int(os.getenv("VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE", "0")) + int(os.getenv("VLLM_SM70_MXFP4_MOE_DIRECT_ORDER_DECODE", "1")) + ), + # Consume the existing TurboMind E2M1/UE8M0 pack directly for the exact + # six-route B1 W13/W2 tensors. Set to 0 to retain the dense-stage path. + "VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE": lambda: bool( + int(os.getenv("VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE", "1")) ), "VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE": lambda: bool( int(os.getenv("VLLM_SM70_MXFP4_MOE_BROADCAST_INPUT_DECODE", "1")) @@ -2622,11 +2698,17 @@ def _resolve_rust_frontend_path() -> str | None: int(os.getenv("VLLM_FLASH_V100_DECODE_USE_XQA", "1")) ), "VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY": lambda: bool( - int(os.getenv("VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY", "0")) + int(os.getenv("VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY", "1")) ), "VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY_MIN_MODEL_LEN": lambda: int( os.getenv("VLLM_FLASH_V100_DFLASH2_GROUPED_VERIFY_MIN_MODEL_LEN", "32768") ), + "VLLM_FLASH_V100_DFLASH2_FIXED_INTERLEAVED": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_DFLASH2_FIXED_INTERLEAVED", "1")) + ), + "VLLM_FLASH_V100_DFLASH2_STAGE_PAGE_IDS": lambda: bool( + int(os.getenv("VLLM_FLASH_V100_DFLASH2_STAGE_PAGE_IDS", "1")) + ), "VLLM_FLASH_V100_DECODE_XQA_Q4_MIN_SEQ_LEN": lambda: int( os.getenv("VLLM_FLASH_V100_DECODE_XQA_Q4_MIN_SEQ_LEN", "32768") ), @@ -2813,6 +2895,36 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_DFLASH_PROFILE_LOG_INTERVAL": lambda: max( 1, int(os.getenv("VLLM_DFLASH_PROFILE_LOG_INTERVAL", "32")) ), + # Lookup-augmented DFlash2 controller. The lookup itself is enabled by the + # speculative config's ngram_assist flag; these knobs tune only its + # experimental fusion and adaptive q8/q16 policy. + "VLLM_DFLASH2_LOOKUP_ADAPTIVE": lambda: bool( + int(os.getenv("VLLM_DFLASH2_LOOKUP_ADAPTIVE", "1")) + ), + "VLLM_DFLASH2_LOOKUP_NSTRONG": lambda: max( + 1, int(os.getenv("VLLM_DFLASH2_LOOKUP_NSTRONG", "6")) + ), + "VLLM_DFLASH2_LOOKUP_AGREE": lambda: max( + 0, int(os.getenv("VLLM_DFLASH2_LOOKUP_AGREE", "0")) + ), + "VLLM_DFLASH2_LOOKUP_NMIN_TAIL": lambda: max( + 1, int(os.getenv("VLLM_DFLASH2_LOOKUP_NMIN_TAIL", "4")) + ), + "VLLM_DFLASH2_LOOKUP_LONG_MIN": lambda: max( + 1, int(os.getenv("VLLM_DFLASH2_LOOKUP_LONG_MIN", "6")) + ), + "VLLM_DFLASH2_LOOKUP_SEARCH": lambda: max( + 1, int(os.getenv("VLLM_DFLASH2_LOOKUP_SEARCH", str(1 << 30))) + ), + "VLLM_DFLASH2_LOOKUP_ENTRY_STREAK": lambda: max( + 1, int(os.getenv("VLLM_DFLASH2_LOOKUP_ENTRY_STREAK", "2")) + ), + "VLLM_DFLASH2_LOOKUP_STICKY": lambda: max( + 0, int(os.getenv("VLLM_DFLASH2_LOOKUP_STICKY", "3")) + ), + "VLLM_DFLASH2_LOOKUP_CHEAP_CONTEXT": lambda: max( + 0, int(os.getenv("VLLM_DFLASH2_LOOKUP_CHEAP_CONTEXT", "0")) + ), "VLLM_DFLASH_DUMP_FIRST_PASS": lambda: bool( int(os.getenv("VLLM_DFLASH_DUMP_FIRST_PASS", "0")) ), @@ -4152,6 +4264,13 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool( int(os.getenv("VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "1")) ), + # Explicit CUDA graph memory reservation for the V2 GPU model runner, + # which does not profile graph memory yet. The default preserves upstream + # behavior. Set this from a measured capture result when sizing a tight KV + # cache pool. + "VLLM_V2_CUDAGRAPH_MEM_MIB": lambda: float( + os.getenv("VLLM_V2_CUDAGRAPH_MEM_MIB", "0") or 0 + ), # NIXL EP environment variables "VLLM_NIXL_EP_MAX_NUM_RANKS": lambda: int( os.getenv("VLLM_NIXL_EP_MAX_NUM_RANKS", "32") diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index e591018d4f..ff900fa467 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -10,6 +10,23 @@ logger = init_logger(__name__) +def _is_exact_sm70_glm_mhc( + residual: torch.Tensor, norm_weight: torch.Tensor | None = None +) -> bool: + if residual.dtype != torch.float16 or not current_platform.is_cuda(): + return False + capability = current_platform.get_device_capability() + if capability is None or capability.to_int() != 70: + return False + if residual.shape[-2:] != (4, 4096): + return False + return norm_weight is None or norm_weight.dtype == torch.float16 + + +def _sm70_mhc_n_splits(num_tokens: int) -> int: + return 8 if num_tokens <= 16 else (4 if num_tokens <= 128 else 2) + + def _require_mhc_activation_dtype( tensor: torch.Tensor, *same_dtype_tensors: tuple[str, torch.Tensor | None], @@ -220,7 +237,12 @@ def mhc_pre_tilelang( from vllm.utils.deep_gemm import is_deep_gemm_supported use_deep_gemm = is_deep_gemm_supported() - if use_deep_gemm: + use_sm70_triton = norm_weight is not None and _is_exact_sm70_glm_mhc( + residual_flat, norm_weight + ) + if use_sm70_triton: + n_splits = _sm70_mhc_n_splits(num_tokens) + elif use_deep_gemm: # these numbers are from deepgemm kernel impl block_k = 64 block_m = 64 @@ -246,7 +268,19 @@ def mhc_pre_tilelang( ) residual_2d = residual_flat.view(num_tokens, hc_mult * hidden_size) - if use_deep_gemm: + if use_sm70_triton: + from vllm.model_executor.kernels.mhc.triton import ( + sm70_mhc_prenorm_staging, + ) + + logger.info_once("SM70 GLM mHC Triton standalone-pre path enabled.") + sm70_mhc_prenorm_staging( + residual_2d, + fn, + gemm_out_mul, + gemm_out_sqrsum, + ) + elif use_deep_gemm: tf32_hc_prenorm_gemm( residual_2d, fn, @@ -264,7 +298,30 @@ def mhc_pre_tilelang( hc_mult, ) - if norm_weight is None: + if use_sm70_triton: + from vllm.model_executor.kernels.mhc.triton import ( + sm70_mhc_pre_norm_from_staging, + ) + + assert norm_weight is not None + sm70_mhc_pre_norm_from_staging( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual_flat, + post_mix, + comb_mix, + layer_input, + norm_weight, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_eps, + ) + elif norm_weight is None: mhc_pre_big_fuse_tilelang( gemm_out_mul, gemm_out_sqrsum, @@ -407,6 +464,16 @@ def mhc_pre_broadcast_tilelang( assert norm_weight is not None num_tokens = residual.shape[0] + capability = ( + current_platform.get_device_capability() if current_platform.is_cuda() else None + ) + use_sm70_triton = ( + use_fp16 + and capability is not None + and capability.to_int() == 70 + and hidden_size == 4096 + and hc_mult == 4 + ) residual_out = torch.empty( num_tokens, hc_mult, @@ -424,7 +491,9 @@ def mhc_pre_broadcast_tilelang( num_tokens, hidden_size, dtype=activation_dtype, device=residual.device ) - if use_fp16: + if use_sm70_triton: + n_splits = 2 if num_tokens <= 16 else 1 + elif use_fp16: # The SM70 path uses the same FP32 accumulation as the regular # TileLang prenorm kernel. Split-K is deliberately disabled: its # V100 shape must satisfy the fixed 512-thread reduction contract. @@ -441,7 +510,20 @@ def mhc_pre_broadcast_tilelang( n_splits, num_tokens, dtype=torch.float32, device=residual.device ) - if use_fp16: + if use_sm70_triton: + from vllm.model_executor.kernels.mhc.triton import ( + sm70_mhc_prenorm_staging, + ) + + logger.info_once("SM70 GLM mHC Triton broadcast-pre path enabled.") + sm70_mhc_prenorm_staging( + residual, + fn_broadcast, + gemm_out_mul, + gemm_out_sqrsum, + sqrsum_mult=hc_mult, + ) + elif use_fp16: _tilelang_hc_prenorm_gemm( residual, fn_broadcast, @@ -462,33 +544,95 @@ def mhc_pre_broadcast_tilelang( n_splits, ) - mhc_pre_big_fuse_broadcast_with_norm_tilelang( - gemm_out_mul, - gemm_out_sqrsum, - hc_scale, - hc_base, - residual, + if use_sm70_triton: + from vllm.model_executor.kernels.mhc.triton import ( + sm70_mhc_pre_norm_from_staging, + ) + + residual_out.copy_(residual.unsqueeze(1)) + sm70_mhc_pre_norm_from_staging( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual_out, + post_mix, + comb_mix, + layer_input, + norm_weight, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_eps, + ) + else: + mhc_pre_big_fuse_broadcast_with_norm_tilelang( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual, + residual_out, + post_mix, + comb_mix, + layer_input, + norm_weight, + hidden_size, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_eps, + n_splits, + hc_mult, + use_fp16=use_fp16, + ) + return ( residual_out, - post_mix, - comb_mix, + post_mix.unsqueeze(-1), + comb_mix.view(num_tokens, hc_mult, hc_mult), layer_input, - norm_weight, - hidden_size, + ) + + +def _mhc_pre_broadcast_tilelang_fake( + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + fn_broadcast: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + del ( + hc_scale, + hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, - norm_eps, n_splits, - hc_mult, - use_fp16=use_fp16, + norm_weight, + norm_eps, + fn_broadcast, ) + num_tokens, hidden_size = residual.shape + hc_mult = fn.shape[1] // hidden_size return ( - residual_out, - post_mix.unsqueeze(-1), - comb_mix.view(num_tokens, hc_mult, hc_mult), - layer_input, + residual.new_empty((num_tokens, hc_mult, hidden_size)), + residual.new_empty((num_tokens, hc_mult, 1), dtype=torch.float32), + residual.new_empty((num_tokens, hc_mult, hc_mult), dtype=torch.float32), + residual.new_empty((num_tokens, hidden_size)), ) @@ -498,11 +642,17 @@ def mhc_post_tilelang( post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor, ) -> torch.Tensor: + activation_dtype = _require_mhc_activation_dtype(residual, ("x", x)) + if _is_exact_sm70_glm_mhc(residual): + from vllm.model_executor.kernels.mhc.triton import sm70_mhc_post + + logger.info_once("SM70 GLM mHC Triton standalone-post path enabled.") + return sm70_mhc_post(x, residual, post_layer_mix, comb_res_mix) + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( mhc_post_tilelang as _mhc_post_kernel, ) - activation_dtype = _require_mhc_activation_dtype(residual, ("x", x)) out = torch.empty_like(residual) _mhc_post_kernel( comb_res_mix, @@ -635,6 +785,13 @@ def mhc_fused_post_pre_tilelang( or (num_tokens == 8 and tile_n == 3 and n_splits == 4) ) ) + use_sm70_triton_large = ( + norm_weight is not None + and _is_exact_sm70_glm_mhc(residual_flat, norm_weight) + and not use_small_fma + ) + if use_sm70_triton_large: + n_splits = _sm70_mhc_n_splits(num_tokens) gemm_out_mul = torch.empty( n_splits, @@ -718,19 +875,42 @@ def mhc_fused_post_pre_tilelang( use_fp16=use_fp16, ) else: - mhc_post_tilelang( - comb_res_mix_flat, - residual_flat, - post_layer_mix_flat, - x_flat, - residual_cur, - residual.shape[-2], - residual.shape[-1], - use_fp16=use_fp16, - ) + if use_sm70_triton_large: + from vllm.model_executor.kernels.mhc.triton import sm70_mhc_post + + logger.info_once("SM70 GLM mHC Triton large-M fused path enabled.") + sm70_mhc_post( + x_flat, + residual_flat, + post_layer_mix_flat, + comb_res_mix_flat, + residual_cur, + ) + else: + mhc_post_tilelang( + comb_res_mix_flat, + residual_flat, + post_layer_mix_flat, + x_flat, + residual_cur, + residual.shape[-2], + residual.shape[-1], + use_fp16=use_fp16, + ) residual_cur_2d = residual_cur.view(num_tokens, hc_mult * hidden_size) - if use_deep_gemm: + if use_sm70_triton_large: + from vllm.model_executor.kernels.mhc.triton import ( + sm70_mhc_prenorm_staging, + ) + + sm70_mhc_prenorm_staging( + residual_cur_2d, + fn, + gemm_out_mul, + gemm_out_sqrsum, + ) + elif use_deep_gemm: from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm tf32_hc_prenorm_gemm( @@ -750,7 +930,37 @@ def mhc_fused_post_pre_tilelang( hc_mult, ) - if norm_weight is None: + use_sm70_pre_norm = ( + use_fp16 + and capability is not None + and capability.to_int() == 70 + and norm_weight is not None + and hidden_size == 4096 + and hc_mult == 4 + ) + if use_sm70_pre_norm: + from vllm.model_executor.kernels.mhc.triton import ( + sm70_mhc_pre_norm_from_staging, + ) + + sm70_mhc_pre_norm_from_staging( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual_cur, + post_mix_cur, + comb_mix_cur, + layer_input_cur, + norm_weight, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_eps, + ) + elif norm_weight is None: mhc_pre_big_fuse_tilelang( gemm_out_mul, gemm_out_sqrsum, @@ -910,6 +1120,12 @@ def _hc_head_fused_kernel_tilelang_fake( mutates_args=[], fake_impl=_mhc_pre_tilelang_fake, ) +direct_register_custom_op( + op_name="mhc_pre_broadcast_tilelang", + op_func=mhc_pre_broadcast_tilelang, + mutates_args=[], + fake_impl=_mhc_pre_broadcast_tilelang_fake, +) direct_register_custom_op( op_name="mhc_post_tilelang", op_func=mhc_post_tilelang, diff --git a/vllm/model_executor/kernels/mhc/torch.py b/vllm/model_executor/kernels/mhc/torch.py index 1a54ce4ad2..a803b02fc3 100644 --- a/vllm/model_executor/kernels/mhc/torch.py +++ b/vllm/model_executor/kernels/mhc/torch.py @@ -14,12 +14,14 @@ def mhc_pre_torch( hc_post_mult_value: float, sinkhorn_repeat: int, n_splits: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Forward pass for mHC pre block. Args: - residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16 + residual: shape (..., hc_mult, hidden_size), dtype bfloat16 or float16 fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32 hc_scale: shape (3,), dtype torch.float32 hc_base: shape (hc_mult3,), dtype torch.float32 @@ -33,11 +35,11 @@ def mhc_pre_torch( Returns: post_mix: shape (..., hc_mult), dtype torch.float32 comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 - layer_input: shape (..., hidden_size), dtype torch.bfloat16 + layer_input: shape (..., hidden_size), dtype matching residual """ # Validate shapes - assert residual.dtype == torch.bfloat16 + assert residual.dtype in (torch.bfloat16, torch.float16) assert fn.dtype == torch.float32 assert hc_scale.dtype == torch.float32 assert hc_base.dtype == torch.float32 @@ -83,7 +85,16 @@ def mhc_pre_torch( layer_input = torch.sum( pre_mix.unsqueeze(-1) * residual_flat.to(torch.float32), dim=1 - ).to(torch.bfloat16) + ) + if norm_weight is not None: + assert norm_weight.shape == (hidden_size,) + variance = layer_input.square().mean(dim=-1, keepdim=True) + layer_input = ( + layer_input + * torch.rsqrt(variance + norm_eps) + * norm_weight.to(torch.float32) + ) + layer_input = layer_input.to(residual.dtype) return ( post_mix.view(*outer_shape, hc_mult, 1), comb_mix.view(*outer_shape, hc_mult, hc_mult), @@ -104,3 +115,40 @@ def mhc_post_torch( ) post_term = post_layer_mix.to(torch.float32) * x.unsqueeze(-2).to(torch.float32) return (mixed_residual + post_term).to(residual.dtype) + + +def mhc_fused_post_pre_torch( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + tile_n: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + del tile_n + residual_cur = mhc_post_torch(x, residual, post_layer_mix, comb_res_mix) + post_mix, comb_mix, layer_input = mhc_pre_torch( + residual_cur, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + norm_weight, + norm_eps, + ) + return residual_cur, post_mix, comb_mix, layer_input diff --git a/vllm/model_executor/kernels/mhc/triton.py b/vllm/model_executor/kernels/mhc/triton.py index facb6813ca..d967a4775d 100644 --- a/vllm/model_executor/kernels/mhc/triton.py +++ b/vllm/model_executor/kernels/mhc/triton.py @@ -4,9 +4,455 @@ import torch.nn.functional as F from torch import Tensor +import vllm._sm70_ops as sm70_ops +from vllm.logger import init_logger from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import direct_register_custom_op +logger = init_logger(__name__) + + +@triton.jit +def _sm70_mhc_sqrsum_staging_kernel( + x_ptr, + sqrsum_ptr, + x_stride_t, + sqrsum_stride_s, + sqrsum_stride_t, + K_PER_SPLIT: tl.constexpr, + SQRSUM_MULT: tl.constexpr, +): + token_idx = tl.program_id(0) + split_idx = tl.program_id(1) + offsets = split_idx * K_PER_SPLIT + tl.arange(0, K_PER_SPLIT) + x = tl.load(x_ptr + token_idx * x_stride_t + offsets).to(tl.float32) + sqrsum = tl.sum(x * x, axis=0) * SQRSUM_MULT + tl.store( + sqrsum_ptr + split_idx * sqrsum_stride_s + token_idx * sqrsum_stride_t, + sqrsum, + ) + + +@triton.jit +def _sm70_mhc_dot_staging_kernel( + x_ptr, + fn_ptr, + out_ptr, + x_stride_t, + fn_stride_n, + out_stride_s, + out_stride_t, + K_PER_SPLIT: tl.constexpr, +): + token_idx = tl.program_id(0) + output_idx = tl.program_id(1) + split_idx = tl.program_id(2) + offsets = split_idx * K_PER_SPLIT + tl.arange(0, K_PER_SPLIT) + x = tl.load(x_ptr + token_idx * x_stride_t + offsets).to(tl.float32) + weight = tl.load(fn_ptr + output_idx * fn_stride_n + offsets).to(tl.float32) + dot = tl.sum(x * weight, axis=0) + tl.store( + out_ptr + split_idx * out_stride_s + token_idx * out_stride_t + output_idx, + dot, + ) + + +def sm70_mhc_prenorm_staging( + x: torch.Tensor, + fn: torch.Tensor, + gemm_out_mul: torch.Tensor, + gemm_out_sqrsum: torch.Tensor, + *, + sqrsum_mult: int = 1, +) -> None: + """Compute exact-FP32 mHC logits and norm partials on V100.""" + if x.dtype != torch.float16 or fn.dtype != torch.float32: + raise TypeError("SM70 mHC prenorm requires FP16 input and FP32 weights") + if x.ndim != 2 or fn.ndim != 2 or not x.is_contiguous(): + raise ValueError("SM70 mHC prenorm requires contiguous 2D tensors") + if fn.shape[1] != x.shape[1] or fn.shape[0] != gemm_out_mul.shape[2]: + raise ValueError("SM70 mHC prenorm staging shapes do not match") + if ( + gemm_out_mul.shape[:2] + != ( + gemm_out_sqrsum.shape[0], + x.shape[0], + ) + or gemm_out_sqrsum.shape[1] != x.shape[0] + ): + raise ValueError("SM70 mHC prenorm output shapes do not match") + n_splits = gemm_out_mul.shape[0] + if x.shape[1] % n_splits != 0: + raise ValueError("SM70 mHC prenorm K dimension must divide n_splits") + k_per_split = x.shape[1] // n_splits + if k_per_split & (k_per_split - 1): + raise ValueError("SM70 mHC prenorm split size must be a power of two") + + _sm70_mhc_sqrsum_staging_kernel[(x.shape[0], n_splits)]( + x, + gemm_out_sqrsum, + x.stride(0), + gemm_out_sqrsum.stride(0), + gemm_out_sqrsum.stride(1), + K_PER_SPLIT=k_per_split, + SQRSUM_MULT=sqrsum_mult, + num_warps=8, + ) + _sm70_mhc_dot_staging_kernel[(x.shape[0], fn.shape[0], n_splits)]( + x, + fn, + gemm_out_mul, + x.stride(0), + fn.stride(0), + gemm_out_mul.stride(0), + gemm_out_mul.stride(1), + K_PER_SPLIT=k_per_split, + num_warps=8, + ) + + +@triton.jit +def _sm70_mhc_post_kernel( + x_ptr, + residual_ptr, + post_mix_ptr, + comb_mix_ptr, + out_ptr, + x_stride_t, + residual_stride_t, + residual_stride_m, + residual_stride_h, + post_stride_t, + post_stride_m, + comb_stride_t, + comb_stride_i, + comb_stride_j, + out_stride_t, + out_stride_m, + out_stride_h, + HIDDEN_SIZE: tl.constexpr, + HC_MULT: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token_idx = tl.program_id(0) + block_idx = tl.program_id(1) + hidden_offsets = block_idx * BLOCK_H + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < HIDDEN_SIZE + x = tl.load( + x_ptr + token_idx * x_stride_t + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + + for output_stream in tl.static_range(0, HC_MULT): + value = ( + tl.load( + post_mix_ptr + token_idx * post_stride_t + output_stream * post_stride_m + ) + * x + ) + for input_stream in tl.static_range(0, HC_MULT): + residual = tl.load( + residual_ptr + + token_idx * residual_stride_t + + input_stream * residual_stride_m + + hidden_offsets * residual_stride_h, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + mix = tl.load( + comb_mix_ptr + + token_idx * comb_stride_t + + input_stream * comb_stride_i + + output_stream * comb_stride_j + ) + value += mix * residual + tl.store( + out_ptr + + token_idx * out_stride_t + + output_stream * out_stride_m + + hidden_offsets * out_stride_h, + value, + mask=hidden_mask, + ) + + +def sm70_mhc_post( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Run the mHC post mapping without TileLang's SM70 BF16 headers.""" + num_tokens, hc_mult, hidden_size = residual.shape + if residual.dtype != torch.float16 or x.dtype != torch.float16: + raise TypeError("SM70 mHC post requires FP16 activations") + if hc_mult != 4 or hidden_size != 4096: + raise ValueError("SM70 mHC post requires hidden_size=4096 and hc_mult=4") + post_layer_mix = post_layer_mix.view(num_tokens, hc_mult) + comb_res_mix = comb_res_mix.view(num_tokens, hc_mult, hc_mult) + if out is None: + out = torch.empty_like(residual) + elif out.shape != residual.shape or out.dtype != residual.dtype: + raise ValueError("SM70 mHC post output must match the residual") + block_h = 512 + _sm70_mhc_post_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( + x.view(num_tokens, hidden_size), + residual, + post_layer_mix, + comb_res_mix, + out, + x.stride(0), + residual.stride(0), + residual.stride(1), + residual.stride(2), + post_layer_mix.stride(0), + post_layer_mix.stride(1), + comb_res_mix.stride(0), + comb_res_mix.stride(1), + comb_res_mix.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + HIDDEN_SIZE=hidden_size, + HC_MULT=hc_mult, + BLOCK_H=block_h, + num_warps=4, + ) + return out + + +@triton.jit +def _sm70_mhc_pre_norm_kernel( + gemm_out_mul_ptr, + gemm_out_sqrsum_ptr, + hc_scale_ptr, + hc_base_ptr, + residual_ptr, + post_mix_ptr, + comb_mix_ptr, + layer_input_ptr, + norm_weight_ptr, + gemm_mul_stride_s, + gemm_mul_stride_t, + gemm_sq_stride_s, + gemm_sq_stride_t, + residual_stride_t, + residual_stride_m, + residual_stride_h, + post_stride_t, + comb_stride_t, + output_stride_t, + rms_eps: tl.constexpr, + hc_pre_eps: tl.constexpr, + hc_sinkhorn_eps: tl.constexpr, + hc_post_mult_value: tl.constexpr, + norm_eps: tl.constexpr, + N_SPLITS: tl.constexpr, + SINKHORN_REPEAT: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + HC_MULT: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """SM70 mHC pre/Sinkhorn/RMSNorm fed by FP32 staging reductions.""" + token_idx = tl.program_id(0) + sqrsum = 0.0 + for split_idx in tl.static_range(0, N_SPLITS): + sqrsum += tl.load( + gemm_out_sqrsum_ptr + + split_idx * gemm_sq_stride_s + + token_idx * gemm_sq_stride_t + ) + input_rsqrt = tl.rsqrt(sqrsum / (HC_MULT * HIDDEN_SIZE) + rms_eps) + + hc_offsets = tl.arange(0, HC_MULT) + post_logits = tl.zeros((HC_MULT,), dtype=tl.float32) + for split_idx in tl.static_range(0, N_SPLITS): + post_logits += tl.load( + gemm_out_mul_ptr + + split_idx * gemm_mul_stride_s + + token_idx * gemm_mul_stride_t + + HC_MULT + + hc_offsets + ) + post_mix = ( + tl.sigmoid( + post_logits * input_rsqrt * tl.load(hc_scale_ptr + 1) + + tl.load(hc_base_ptr + HC_MULT + hc_offsets) + ) + * hc_post_mult_value + ) + tl.store(post_mix_ptr + token_idx * post_stride_t + hc_offsets, post_mix) + + comb_offsets = tl.arange(0, 16) + comb_logits = tl.zeros((16,), dtype=tl.float32) + for split_idx in tl.static_range(0, N_SPLITS): + comb_logits += tl.load( + gemm_out_mul_ptr + + split_idx * gemm_mul_stride_s + + token_idx * gemm_mul_stride_t + + 2 * HC_MULT + + comb_offsets + ) + comb_logits = comb_logits * input_rsqrt * tl.load(hc_scale_ptr + 2) + tl.load( + hc_base_ptr + 2 * HC_MULT + comb_offsets + ) + comb = tl.zeros((16,), dtype=tl.float32) + for row_idx in tl.static_range(0, HC_MULT): + row_mask = comb_offsets // HC_MULT == row_idx + row_max = tl.max(tl.where(row_mask, comb_logits, -float("inf")), axis=0) + row_exp = tl.exp(comb_logits - row_max) + row_sum = tl.sum(tl.where(row_mask, row_exp, 0.0), axis=0) + comb = tl.where(row_mask, row_exp / row_sum + hc_sinkhorn_eps, comb) + + for col_idx in tl.static_range(0, HC_MULT): + col_mask = comb_offsets % HC_MULT == col_idx + col_sum = tl.sum(tl.where(col_mask, comb, 0.0), axis=0) + comb = tl.where(col_mask, comb / (col_sum + hc_sinkhorn_eps), comb) + + for _ in tl.static_range(0, SINKHORN_REPEAT - 1): + for row_idx in tl.static_range(0, HC_MULT): + row_mask = comb_offsets // HC_MULT == row_idx + row_sum = tl.sum(tl.where(row_mask, comb, 0.0), axis=0) + comb = tl.where(row_mask, comb / (row_sum + hc_sinkhorn_eps), comb) + for col_idx in tl.static_range(0, HC_MULT): + col_mask = comb_offsets % HC_MULT == col_idx + col_sum = tl.sum(tl.where(col_mask, comb, 0.0), axis=0) + comb = tl.where(col_mask, comb / (col_sum + hc_sinkhorn_eps), comb) + tl.store(comb_mix_ptr + token_idx * comb_stride_t + comb_offsets, comb) + + hidden_offsets = tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < HIDDEN_SIZE + layer_input = tl.zeros((BLOCK_H,), dtype=tl.float32) + for stream_idx in tl.static_range(0, HC_MULT): + pre_logit = 0.0 + for split_idx in tl.static_range(0, N_SPLITS): + pre_logit += tl.load( + gemm_out_mul_ptr + + split_idx * gemm_mul_stride_s + + token_idx * gemm_mul_stride_t + + stream_idx + ) + pre_mix = ( + tl.sigmoid( + pre_logit * input_rsqrt * tl.load(hc_scale_ptr) + + tl.load(hc_base_ptr + stream_idx) + ) + + hc_pre_eps + ) + residual = tl.load( + residual_ptr + + token_idx * residual_stride_t + + stream_idx * residual_stride_m + + hidden_offsets * residual_stride_h, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + layer_input += pre_mix * residual + + variance = tl.sum(layer_input * layer_input, axis=0) / HIDDEN_SIZE + norm_scale = tl.rsqrt(variance + norm_eps) + norm_weight = tl.load( + norm_weight_ptr + hidden_offsets, mask=hidden_mask, other=0.0 + ).to(tl.float32) + # Match the upstream fused kernel's FP16 staging before RMSNorm scaling. + layer_input = layer_input.to(tl.float16).to(tl.float32) + layer_input = layer_input * norm_scale * norm_weight + tl.store( + layer_input_ptr + token_idx * output_stride_t + hidden_offsets, + layer_input, + mask=hidden_mask, + ) + + +def sm70_mhc_pre_norm_from_staging( + gemm_out_mul: torch.Tensor, + gemm_out_sqrsum: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + layer_input: torch.Tensor, + norm_weight: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + norm_eps: float, +) -> None: + """Launch the SM70 mHC final stage without BF16 codegen.""" + num_tokens, hc_mult, hidden_size = residual.shape + if hidden_size != 4096 or hc_mult != 4: + raise ValueError( + "SM70 mHC staging kernel requires hidden_size=4096 and hc_mult=4" + ) + if num_tokens < 1: + raise ValueError("SM70 mHC staging kernel requires at least one token") + if norm_weight.dtype != torch.float16: + raise TypeError("SM70 mHC staging kernel requires FP16 norm weights") + + if num_tokens == 1: + if not hasattr(torch.ops._C, "sm70_glm_mhc_pre_norm_out"): + raise RuntimeError( + "SM70 GLM mHC decode requires the native CUDA final-stage op. " + "Rebuild vLLM from source with CUDA arch 7.0." + ) + logger.info_once("SM70 GLM mHC native CUDA decode final stage enabled.") + sm70_ops.sm70_glm_mhc_pre_norm_out( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual, + post_mix, + comb_mix, + layer_input, + norm_weight, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_eps, + ) + return + + _sm70_mhc_pre_norm_kernel[(num_tokens,)]( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + residual, + post_mix, + comb_mix, + layer_input, + norm_weight, + gemm_out_mul.stride(0), + gemm_out_mul.stride(1), + gemm_out_sqrsum.stride(0), + gemm_out_sqrsum.stride(1), + residual.stride(0), + residual.stride(1), + residual.stride(2), + post_mix.stride(0), + comb_mix.stride(0), + layer_input.stride(0), + rms_eps=rms_eps, + hc_pre_eps=hc_pre_eps, + hc_sinkhorn_eps=hc_sinkhorn_eps, + hc_post_mult_value=hc_post_mult_value, + norm_eps=norm_eps, + N_SPLITS=gemm_out_mul.shape[0], + SINKHORN_REPEAT=sinkhorn_repeat, + HIDDEN_SIZE=hidden_size, + HC_MULT=hc_mult, + BLOCK_H=triton.next_power_of_2(hidden_size), + num_warps=8, + ) + @triton.jit def _rmsnorm_nw_kernel( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 71fd297a7e..0c42f89715 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -669,7 +669,11 @@ def forward_impl( k_c_normed = k_c_normed[:num_actual_toks, ...] k_pe = k_pe[:num_actual_toks, ...] - if fp8_attention and self.kv_cache_dtype != "fp8_ds_mla": + if ( + fp8_attention + and self.kv_cache_dtype != "fp8_ds_mla" + and self.attn_backend.get_name() != "GLM5_SM70_SPARSE" + ): kv_cache = kv_cache.view(current_platform.fp8_dtype()) # Sparse MLA impls only support forward_mqa (decode-style attention) @@ -961,12 +965,19 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: kv_cache_dtype = kv_cache_dtype_str_to_dtype( self.kv_cache_dtype, vllm_config.model_config ) + model_version = ( + "glm5_next" + if self.attn_backend.get_name() == "GLM5_SM70_SPARSE" + and self.kv_cache_dtype in {"fp8", "fp8_e4m3"} + else None + ) return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=1, head_size=self.head_size, dtype=kv_cache_dtype, cache_dtype_str=vllm_config.cache_config.cache_dtype, + model_version=model_version, ) def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor): diff --git a/vllm/model_executor/layers/attention_layer_base.py b/vllm/model_executor/layers/attention_layer_base.py index 97395b6414..512ffd182c 100644 --- a/vllm/model_executor/layers/attention_layer_base.py +++ b/vllm/model_executor/layers/attention_layer_base.py @@ -4,6 +4,8 @@ from abc import ABC, abstractmethod +import torch + from vllm.config import VllmConfig from vllm.v1.attention.backend import AttentionBackend, AttentionImpl from vllm.v1.kv_cache_interface import KVCacheSpec @@ -20,6 +22,10 @@ class AttentionLayerBase(ABC): impl: "AttentionImpl" + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + """Bind an allocated KV cache view to this layer.""" + self.kv_cache = kv_cache + @abstractmethod def get_attn_backend(self) -> type[AttentionBackend]: """Get the attention backend class for this layer.""" diff --git a/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py b/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py index d0735721da..f86223b07c 100644 --- a/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py +++ b/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py @@ -45,7 +45,9 @@ def _parse_int_list(env_name: str, default_vals: list[int]) -> list[int]: return out or default_vals -_use_sm70_kkt_schedule = os.getenv("VLLM_SM70_GDN_KKT_SCHEDULE", "1") == "1" and _is_sm70() +_use_sm70_kkt_schedule = ( + os.getenv("VLLM_SM70_GDN_KKT_SCHEDULE", "1") == "1" and _is_sm70() +) _kkt_configs = ( [ triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) diff --git a/vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py b/vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py index 87708d2675..fecfa3a4e7 100644 --- a/vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py +++ b/vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py @@ -76,8 +76,7 @@ def _fused_post_conv_kernel( # Load Q features: mixed_qkv[t, i_h*K + k] q_offsets = ( - offs_t[:, None] * stride_x_tok - + (i_h * K + offs_k[None, :]) * stride_x_feat + offs_t[:, None] * stride_x_tok + (i_h * K + offs_k[None, :]) * stride_x_feat ) q_f32 = tl.load(mixed_qkv_ptr + q_offsets, mask=qk_mask_2d, other=0).to( tl.float32 diff --git a/vllm/model_executor/layers/fla/ops/fused_recurrent.py b/vllm/model_executor/layers/fla/ops/fused_recurrent.py index ed1d1a1ef4..4e47bd5e49 100644 --- a/vllm/model_executor/layers/fla/ops/fused_recurrent.py +++ b/vllm/model_executor/layers/fla/ops/fused_recurrent.py @@ -56,9 +56,7 @@ def _round_num_warps(value: int) -> int: return 8 -_SM70_FLA_RECURRENT_SCHEDULE = ( - os.getenv("VLLM_SM70_FLA_RECURRENT_SCHEDULE", "1") == "1" -) +_SM70_FLA_RECURRENT_SCHEDULE = os.getenv("VLLM_SM70_FLA_RECURRENT_SCHEDULE", "1") == "1" _SM70_FLA_BV_OVERRIDE = _parse_positive_int_env("VLLM_SM70_FLA_BV") _SM70_FLA_WARPS_OVERRIDE = _parse_positive_int_env("VLLM_SM70_FLA_WARPS") _SM70_FLA_STAGES_OVERRIDE = _parse_positive_int_env("VLLM_SM70_FLA_STAGES") @@ -154,6 +152,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( cu_seqlens, ssm_state_indices, num_accepted_tokens, + a_log, + g_bias, scale, N: tl.int64, # num of sequences T: tl.int64, # num of tokens @@ -176,6 +176,10 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( IS_CONTINUOUS_BATCHING: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, IS_KDA: tl.constexpr, + SIGMOID_BETA: tl.constexpr, + COMPUTE_GATE: tl.constexpr, + SAFE_GATE: tl.constexpr, + LOWER_BOUND: tl.constexpr, ): i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_n, i_hv = i_nh // HV, i_nh % HV @@ -211,6 +215,9 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( else: p_gk = g + (bos * HV + i_hv) * K + o_k + if COMPUTE_GATE: + b_a_log = tl.exp(tl.load(a_log + i_h).to(tl.float32)) + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v mask_k = o_k < K @@ -253,6 +260,12 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_h *= exp(b_g) else: b_gk = tl.load(p_gk).to(tl.float32) + if COMPUTE_GATE: + b_gk += tl.load(g_bias + i_h * K + o_k, mask=mask_k, other=0.0).to( + tl.float32 + ) + if SAFE_GATE: + b_gk = LOWER_BOUND / (1.0 + tl.exp(-(b_a_log * b_gk))) b_h *= exp(b_gk[None, :]) # [BV] b_v -= tl.sum(b_h * b_k[None, :], 1) @@ -260,6 +273,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) else: b_beta = tl.load(p_beta).to(tl.float32) + if SIGMOID_BETA: + b_beta = tl.sigmoid(b_beta) b_v *= b_beta # [BV, BK] b_h += b_v[:, None] * b_k[None, :] @@ -351,6 +366,8 @@ def fused_recurrent_gated_delta_rule_fwd( cu_seqlens=cu_seqlens, ssm_state_indices=ssm_state_indices, num_accepted_tokens=num_accepted_tokens, + a_log=None, + g_bias=None, scale=scale, N=N, T=T, @@ -369,6 +386,10 @@ def fused_recurrent_gated_delta_rule_fwd( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, INPLACE_FINAL_STATE=inplace_final_state, IS_KDA=False, + SIGMOID_BETA=False, + COMPUTE_GATE=False, + SAFE_GATE=True, + LOWER_BOUND=-5.0, num_warps=num_warps, num_stages=num_stages, ) diff --git a/vllm/model_executor/layers/fla/ops/kda.py b/vllm/model_executor/layers/fla/ops/kda.py index 10acc214f4..1eafa490c1 100644 --- a/vllm/model_executor/layers/fla/ops/kda.py +++ b/vllm/model_executor/layers/fla/ops/kda.py @@ -42,6 +42,12 @@ def fused_recurrent_kda_fwd( ssm_state_indices: torch.Tensor | None = None, num_accepted_tokens: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, + out: torch.Tensor | None = None, + sigmoid_beta: bool = False, + a_log: torch.Tensor | None = None, + g_bias: torch.Tensor | None = None, + compute_gate: bool = False, + lower_bound: float | None = -5.0, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] HV = v.shape[2] @@ -52,7 +58,24 @@ def fused_recurrent_kda_fwd( num_stages = 3 num_warps = 1 - o = torch.empty_like(k) + if compute_gate: + assert a_log is not None and g_bias is not None, ( + "compute_gate requires a_log and g_bias" + ) + assert lower_bound is not None, ( + "compute_gate implements the bounded (safe_gate) branch only" + ) + a_log = a_log.reshape(-1).contiguous() + g_bias = g_bias.reshape(-1).contiguous() + + if out is None: + o = torch.empty_like(k) + else: + # Caller-provided output buffer; must be layout-compatible with the + # tensor the kernel indexes (contiguous, same shape/dtype as k). + assert out.shape == k.shape and out.dtype == k.dtype + assert out.is_contiguous() + o = out if inplace_final_state: final_state = initial_state else: @@ -99,6 +122,12 @@ def fused_recurrent_kda_fwd( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, INPLACE_FINAL_STATE=inplace_final_state, IS_KDA=True, + SIGMOID_BETA=sigmoid_beta, + a_log=a_log, + g_bias=g_bias, + COMPUTE_GATE=compute_gate, + SAFE_GATE=True, + LOWER_BOUND=lower_bound if lower_bound is not None else -5.0, num_warps=num_warps, num_stages=num_stages, ) @@ -118,6 +147,13 @@ def fused_recurrent_kda( use_qk_l2norm_in_kernel: bool = True, cu_seqlens: torch.Tensor | None = None, ssm_state_indices: torch.LongTensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + out: torch.Tensor | None = None, + sigmoid_beta: bool = False, + a_log: torch.Tensor | None = None, + g_bias: torch.Tensor | None = None, + compute_gate: bool = False, + lower_bound: float | None = -5.0, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: if cu_seqlens is not None and q.shape[0] != 1: @@ -139,8 +175,14 @@ def fused_recurrent_kda( inplace_final_state=inplace_final_state, cu_seqlens=cu_seqlens, ssm_state_indices=ssm_state_indices, - num_accepted_tokens=None, + num_accepted_tokens=num_accepted_tokens, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + out=out, + sigmoid_beta=sigmoid_beta, + a_log=a_log, + g_bias=g_bias, + compute_gate=compute_gate, + lower_bound=lower_bound, ) return o, final_state @@ -1189,6 +1231,8 @@ def kda_gate_cumsum_fwd_kernel( cumsum_scale, beta, threshold, + SAFE_GATE: tl.constexpr, + LOWER_BOUND: tl.constexpr, T, H: tl.constexpr, D: tl.constexpr, @@ -1235,14 +1279,20 @@ def kda_gate_cumsum_fwd_kernel( b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] - b_a = -tl.exp(tl.load(A + i_h).to(tl.float32)) - b_g_scaled = b_g * beta - b_softplus = tl.where( - b_g_scaled > threshold, - b_g, - (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), - ) - b_gate = b_a * b_softplus + b_a = tl.load(A + i_h).to(tl.float32) + b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a) + if SAFE_GATE: + # y = lower_bound * sigmoid(exp(A) * (g + g_bias)), bounded to + # (lower_bound, 0) for safe-gate checkpoints. + b_gate = LOWER_BOUND / (1.0 + tl.exp(-(b_a * b_g))) + else: + b_g_scaled = b_g * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_g, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = b_a * b_softplus # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to @@ -1264,6 +1314,8 @@ def fused_kda_gate_chunk_cumsum( chunk_indices: torch.Tensor | None = None, chunk_size: int = FLA_CHUNK_SIZE, output_dtype: torch.dtype | None = torch.float, + safe_gate: bool = False, + lower_bound: float = -5.0, ) -> torch.Tensor: if cu_seqlens is not None: assert raw_g.shape[0] == 1, ( @@ -1295,6 +1347,8 @@ def grid(meta): cumsum_scale=RCP_LN2, beta=beta, threshold=threshold, + SAFE_GATE=safe_gate, + LOWER_BOUND=lower_bound, T=T, H=H, D=D, @@ -1425,6 +1479,8 @@ def chunk_kda_with_fused_gate_fwd( initial_state: torch.Tensor, output_final_state: bool, cu_seqlens: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float = -5.0, ): chunk_size = FLA_CHUNK_SIZE chunk_indices = ( @@ -1439,6 +1495,8 @@ def chunk_kda_with_fused_gate_fwd( cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, ) return _chunk_kda_fwd_with_cumulative_g( q=q, @@ -1502,6 +1560,8 @@ def chunk_kda_with_fused_gate( output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float = -5.0, **kwargs, ): """Run chunk KDA from raw gate projection using fused gate+cumsum.""" @@ -1524,6 +1584,8 @@ def chunk_kda_with_fused_gate( initial_state=initial_state.contiguous() if initial_state is not None else None, output_final_state=output_final_state, cu_seqlens=cu_seqlens, + safe_gate=safe_gate, + lower_bound=lower_bound, ) return o, final_state @@ -1545,6 +1607,8 @@ def kda_gate_fwd_kernel( g_bias, beta: tl.constexpr, threshold: tl.constexpr, + SAFE_GATE: tl.constexpr, + LOWER_BOUND: tl.constexpr, T, H, D: tl.constexpr, @@ -1556,7 +1620,7 @@ def kda_gate_fwd_kernel( n_t = i_t * BT b_a = tl.load(A + i_h).to(tl.float32) - b_a = -tl.exp(b_a) + b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a) stride_row = H * D stride_col = 1 @@ -1589,13 +1653,18 @@ def kda_gate_fwd_kernel( ) b_g = b_g + b_bias[None, :] - # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x)) - # When beta * x > threshold, use linear approximation x - # Use threshold to switch to linear when beta*x > threshold - g_scaled = b_g * beta - use_linear = g_scaled > threshold - sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) - b_y = b_a * sp + if SAFE_GATE: + # y = lower_bound * sigmoid(exp(A) * (g + g_bias)), bounded to + # (lower_bound, 0) for safe-gate checkpoints. + b_y = LOWER_BOUND / (1.0 + tl.exp(-(b_a * b_g))) + else: + # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x)) + # When beta * x > threshold, use linear approximation x + # Use threshold to switch to linear when beta*x > threshold + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = b_a * sp tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) @@ -1607,13 +1676,18 @@ def fused_kda_gate( g_bias: torch.Tensor | None = None, beta: float = 1.0, threshold: float = 20.0, + safe_gate: bool = False, + lower_bound: float | None = -5.0, ) -> torch.Tensor: """ Forward pass for KDA gate: input g: [..., H*D] param A: [H] or [1, 1, H, 1] - beta: softplus beta parameter - threshold: softplus threshold parameter + beta: softplus beta parameter (softplus branch only) + threshold: softplus threshold parameter (softplus branch only) + safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias); + when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias)) + lower_bound: floor for the safe_gate branch (default -5.0) return : [..., H, D] """ orig_shape = g.shape[:-1] @@ -1636,6 +1710,8 @@ def grid(meta): g_bias, beta, threshold, + safe_gate, + lower_bound if lower_bound is not None else -5.0, T, H, head_k_dim, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 4ff43ce21b..beb1f6be7c 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -1338,6 +1338,7 @@ def make_expert_params_mapping( ckpt_up_proj_name: str, num_experts: int, num_redundant_experts: int = 0, + include_fused: bool = False, ) -> list[tuple[str, str, int, str]]: num_physical_experts = num_experts + num_redundant_experts @@ -1357,7 +1358,48 @@ def make_expert_params_mapping( else "" ) - return [ + fused_mapping: list[tuple[str, str, int, str]] = [] + if include_fused: + gate_up_name = None + if ckpt_gate_proj_name == "gate_proj" and ckpt_up_proj_name == "up_proj": + gate_up_name = "gate_up_proj" + elif ckpt_gate_proj_name == "w1" and ckpt_up_proj_name == "w3": + gate_up_name = "w13" + else: + logger.warning( + "Unexpected gate/up projection names: %s, %s. " + "Fused gate/up mapping will be skipped.", + ckpt_gate_proj_name, + ckpt_up_proj_name, + ) + + if gate_up_name is not None: + # Some checkpoints store every expert in one 3D tensor and + # concatenate gate/up along dim 1. FusedMoE.load_weights + # already knows how to split and TP-shard that layout; these + # aliases make the checkpoint names reachable. + fused_mapping = [ + ( + f"experts.{base_layer}w13_weight", + f"experts.{gate_up_name}", + 0, + "w1", + ), + ( + f"experts.{base_layer}w13_weight", + f"experts.{gate_up_name}", + 1, + "w3", + ), + ( + f"experts.{base_layer}w2_weight", + f"experts.{ckpt_down_proj_name}", + 0, + "w2", + ), + ] + + per_expert_mapping = [ # (param_name, weight_name, expert_id, shard_id) ( f"experts.{base_layer}w13_" @@ -1374,6 +1416,7 @@ def make_expert_params_mapping( ("w3", ckpt_up_proj_name), ] ] + return fused_mapping + per_expert_mapping @property def hidden_size(self) -> int: @@ -1404,6 +1447,7 @@ def fused_moe_make_expert_params_mapping( ckpt_up_proj_name: str, num_experts: int, num_redundant_experts: int = 0, + include_fused: bool = False, ) -> list[tuple[str, str, int, str]]: return FusedMoE.make_expert_params_mapping( model, @@ -1412,6 +1456,7 @@ def fused_moe_make_expert_params_mapping( ckpt_up_proj_name, num_experts, num_redundant_experts, + include_fused, ) diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py index a4800eabb9..6aee9a53eb 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py @@ -5,13 +5,94 @@ import torch import vllm._custom_ops as ops +from vllm import envs from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, get_routing_method_type, ) from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +logger = init_logger(__name__) + + +@triton.jit +def _sm70_qwen38_router_topk_kernel( + gating_ptr, + topk_weights_ptr, + topk_ids_ptr, + token_expert_indices_ptr, + E: tl.constexpr, + K: tl.constexpr, + BLOCK_E: tl.constexpr, +) -> None: + """Sort the exact Qwen3.8 decode router in one SM70 program.""" + + offsets = tl.arange(0, BLOCK_E) + valid = offsets < E + logits = tl.load(gating_ptr + offsets, mask=valid, other=-float("inf")).to( + tl.float32 + ) + max_logit = tl.max(logits, axis=0) + + # Match topk_softmax's degenerate-row behavior: NaN, +Inf, or all -Inf + # produces the first K expert IDs with zero weights. + has_nan = tl.max((logits != logits).to(tl.int32), axis=0) != 0 + invalid_row = has_nan | (max_logit == float("inf")) | (max_logit == -float("inf")) + sort_logits = tl.where(invalid_row, -offsets.to(tl.float32), logits) + # Numeric ties use the lower expert ID in the generic CUDA op. Canonicalize + # signed zero before bit packing so -0.0 and +0.0 remain one tie class. + sort_logits = tl.where(sort_logits == 0.0, 0.0, sort_logits) + + # Transform float32 into an ascending-sortable key. Packing the expert ID + # into the low bits preserves the generic kernel's lower-ID tie break. + min_i32: tl.constexpr = -2147483648 + logit_bits = sort_logits.to(tl.int32, bitcast=True) + sign = logit_bits >> 31 + key = tl.where(sign == 0, logit_bits ^ -1, logit_bits ^ min_i32) + key = tl.where(valid, key, 0x7FFFFFFF) + packed = ((key.to(tl.int64) & 0xFFFFFFFF) << 32) | offsets.to(tl.int64) + sorted_packed = tl.sort(packed, descending=False) + + sorted_keys = ((sorted_packed >> 32) & 0xFFFFFFFF).to(tl.int32) + sorted_ids = (sorted_packed & 0xFFFFFFFF).to(tl.int32) + sorted_sign = sorted_keys >> 31 + sorted_bits = tl.where(sorted_sign < 0, sorted_keys ^ -1, sorted_keys ^ min_i32) + sorted_logits = sorted_bits.to(tl.float32, bitcast=True) + + raw_weights = tl.math.exp2((sorted_logits - max_logit) * 1.4426950408889634) + raw_weights = tl.where(invalid_row, 0.0, raw_weights) + top_mask = offsets < K + denominator = tl.sum(tl.where(top_mask, raw_weights, 0.0), axis=0) + denominator = tl.where(denominator > 0.0, denominator, 1.0) + weights = raw_weights / denominator + + tl.store(topk_ids_ptr + offsets, sorted_ids, mask=top_mask) + tl.store(topk_weights_ptr + offsets, weights, mask=top_mask) + # The generic op stores rank-major source rows. M is exactly one here. + tl.store(token_expert_indices_ptr + offsets, offsets, mask=top_mask) + + +def _sm70_qwen38_router_topk( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, +) -> None: + _sm70_qwen38_router_topk_kernel[(1,)]( + gating_output, + topk_weights, + topk_ids, + token_expert_indices, + E=512, + K=10, + BLOCK_E=512, + num_warps=8, + ) def vllm_topk_softmax( @@ -92,6 +173,26 @@ def fused_topk( ) if scoring_func == "softmax": + if ( + envs.VLLM_SM70_QWEN38_ROUTER_TOPK + and M == 1 + and gating_output.shape == (1, 512) + and gating_output.dtype == torch.float16 + and gating_output.is_contiguous() + and topk == 10 + and renormalize + and topk_ids.dtype == torch.int32 + and current_platform.is_device_capability(70) + ): + logger.info_once("SM70 Qwen3.8 E512/K10 router top-k path enabled.") + _sm70_qwen38_router_topk( + topk_weights, + topk_ids, + token_expert_indices, + gating_output, + ) + return topk_weights, topk_ids, token_expert_indices + topk_func = dispatch_topk_softmax_func( use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled() ) diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 414bc561d1..4b485202ad 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -57,6 +57,7 @@ def __init__( prefix=prefix, ) self.out_dtype = out_dtype + self._sm70_dsv4_fp13_gemv = True # DSV3 specialized kernel eligibility (SM90+, exact dims) self.allow_specialized_router_gemm = can_use_specialized_kernels @@ -95,12 +96,19 @@ def forward( ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: import vllm._custom_ops as ops - if envs.VLLM_SM70_DSV4_FP16_GEMV and self.out_dtype is not None: + if ( + envs.VLLM_SM70_DSV4_FP16_GEMV or envs.VLLM_SM70_DSV4_FP13_GEMV + ) and self.out_dtype is not None: from vllm.models.deepseek_v4.sm70.gemv import ( maybe_sm70_dsv4_fp16_gemv, ) - output = maybe_sm70_dsv4_fp16_gemv(x, self.weight, self.out_dtype) + output = maybe_sm70_dsv4_fp16_gemv( + x, + self.weight, + self.out_dtype, + getattr(self, "_sm70_dsv4_fp13_weight", None), + ) if output is not None: return output, None diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index e84990c564..19f74ab5f7 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -3,6 +3,7 @@ import itertools from abc import abstractmethod +from collections.abc import Iterable import torch from torch.nn.parameter import Parameter, UninitializedParameter @@ -369,6 +370,22 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if not current_platform.is_cuda_alike(): return + from vllm.model_executor.layers.quantization.sm70_online_qpn8 import ( + maybe_prepare_online_qpn8, + ) + + if maybe_prepare_online_qpn8(layer): + return + + if envs.VLLM_SM70_DSV4_FP13_GEMV and getattr( + layer, "_sm70_dsv4_fp13_gemv", False + ): + from vllm.models.deepseek_v4.sm70.gemv import ( + prepare_sm70_dsv4_fp13_gemv, + ) + + prepare_sm70_dsv4_fp13_gemv(layer) + force_enable = getattr(layer, "_sm70_f16_force_enable", False) if not envs.VLLM_SM70_ENABLE_DENSE_F16_FASTPATH and not force_enable: return @@ -425,6 +442,13 @@ def apply( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + from vllm.model_executor.layers.quantization.sm70_online_qpn8 import ( + maybe_apply_online_qpn8, + ) + + online_qpn8_out = maybe_apply_online_qpn8(layer, x, bias) + if online_qpn8_out is not None: + return online_qpn8_out sm70_out = _maybe_sm70_dense_forward(layer, x, bias) if sm70_out is not None: return sm70_out @@ -1206,6 +1230,31 @@ def weight_loader_v2( tp_rank=self.tp_rank, ) + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[str]: + for name, loaded_weight in weights: + shard_id = getattr(loaded_weight, "shard_id", None) + self.validate_shard_id(shard_id) + if "." in name: + submodule, _, attr = name.rpartition(".") + param = getattr(self.get_submodule(submodule), attr, None) + else: + param = getattr(self, name, None) + if param is None and name == "bias": + continue + if param is None: + raise ValueError(f"Unknown parameter {name!r} for {self.prefix}") + param.weight_loader(param, loaded_weight, shard_id) + logger.debug( + "Loaded shard %s with shape %s into %s.%s", + shard_id, + loaded_weight.shape, + self.prefix, + name, + ) + yield name + class QKVParallelLinear(ColumnParallelLinear): """Linear layers for the attention's QKV transformation. @@ -1304,6 +1353,32 @@ def validate_shard_id(self, loaded_shard_id: str | None): return raise ValueError("This line should not be reached") + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[str]: + """Load separately serialized Q/K/V shards through AutoWeightsLoader.""" + for name, loaded_weight in weights: + shard_id = getattr(loaded_weight, "shard_id", None) + self.validate_shard_id(shard_id) + if "." in name: + submodule, _, attr = name.rpartition(".") + param = getattr(self.get_submodule(submodule), attr, None) + else: + param = getattr(self, name, None) + if param is None and name == "bias": + continue + if param is None: + raise ValueError(f"Unknown parameter {name!r} for {self.prefix}") + param.weight_loader(param, loaded_weight, shard_id) + logger.debug( + "Loaded shard %s with shape %s into %s.%s", + shard_id, + loaded_weight.shape, + self.prefix, + name, + ) + yield name + def _get_shard_offset_mapping(self, loaded_shard_id: str): shard_offset_mapping = { "q": 0, diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index f8b88a5e8d..13aaf64b3a 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -37,6 +37,10 @@ def get_state_shape(self) -> Iterable[tuple[int, ...]]: def mamba_type(self) -> MambaAttentionBackendEnum: pass + @property + def is_kv_cache_tp_replicated(self) -> bool: + return False + @abstractmethod def get_state_dtype(self) -> tuple[torch.dtype, ...]: pass @@ -51,6 +55,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: block_size=mamba_block_size, page_size_padded=page_size_padded, mamba_type=self.mamba_type, + tp_replicated=self.is_kv_cache_tp_replicated, mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode, num_speculative_blocks=( vllm_config.speculative_config.num_speculative_state_tokens() diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 5f7c5ce0bf..e5d8ba03ef 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -105,6 +105,29 @@ _DFLASH_DDTREE_PATH_PROBE_REPORTS = 0 +def _warmup_sm70_qwen_gdn_causal_conv1d( + forward_context: dict[str, object], +) -> bool: + """Warm one bound Qwen GDN layer using its production cache layout. + + MRV2 owns the cache on the modules registered in ``forward_context``. + Scanning ``model.modules()`` therefore misses the bound cache and leaves + the non-speculative causal-conv variant to JIT on the first structured + request. + """ + if ( + not envs.VLLM_SM70_AUX_KERNEL_WARMUP + or not current_platform.is_device_capability(70) + ): + return False + + for layer in forward_context.values(): + warmup = getattr(layer, "_warmup_sm70_causal_conv1d_real_state", None) + if warmup is not None and warmup(): + return True + return False + + @triton.jit def _sm70_pack_qwen_gdn_qkv_kernel( mixed_qkv, @@ -2374,7 +2397,7 @@ def __init__( _sm70_gdn_rmsnorm_onepass_enabled() and current_platform.is_device_capability(70) and self._sm70_spec_cache_stride == 1 - and self.hidden_size == 5120 + and self.hidden_size in (2560, 5120) and self.tp_size == 4 and self.num_v_heads == 48 and self.head_v_dim == 128 @@ -2383,7 +2406,7 @@ def __init__( self.enable_sm70_gdn_qpn8_ba_split = ( current_platform.is_device_capability(70) and self._sm70_spec_cache_stride == 1 - and self.hidden_size == 5120 + and self.hidden_size in (2560, 5120) and self.tp_size == 4 and self.num_v_heads == 48 and self.head_v_dim == 128 @@ -7127,6 +7150,9 @@ def _sm70_gdn_qpn8_ba_weight_contract( qkvz_weight = getattr(qkvz, "weight", None) qkvz_scales = getattr(qkvz, "weight_scale_inv", None) ba_weight = getattr(ba, "weight", None) if ba is not None else None + input_width = ( + int(qkvz_weight.shape[0]) if isinstance(qkvz_weight, torch.Tensor) else 0 + ) return bool( self.enable_sm70_gdn_qpn8_ba_split and ba is not None @@ -7134,7 +7160,8 @@ def _sm70_gdn_qpn8_ba_weight_contract( and int(getattr(qkvz, "sm70_fp8_prefill_exact_dense_workspace_ptr", 0)) > 0 and isinstance(qkvz_weight, torch.Tensor) and qkvz_weight.dtype == torch.uint8 - and qkvz_weight.shape == (5120, 4096) + and input_width in (2560, 5120) + and qkvz_weight.shape == (input_width, 4096) and qkvz_weight.is_contiguous() and isinstance(qkvz_scales, torch.Tensor) and qkvz_scales.dtype == torch.float16 @@ -7142,7 +7169,7 @@ def _sm70_gdn_qpn8_ba_weight_contract( and qkvz_scales.is_contiguous() and isinstance(ba_weight, torch.Tensor) and ba_weight.dtype == torch.float16 - and ba_weight.shape == (24, 5120) + and ba_weight.shape == (24, input_width) and ba_weight.is_contiguous() and qkvz_weight.device == qkvz_scales.device == ba_weight.device and getattr(qkvz, "bias", None) is None @@ -7166,7 +7193,7 @@ def _sm70_gdn_qpn8_ba_dispatch_eligible( and hidden_states.is_cuda and hidden_states.dtype == torch.float16 and hidden_states.ndim == 2 - and hidden_states.shape[1] == 5120 + and hidden_states.shape[1] == self.in_proj_qkvz.weight.shape[0] and hidden_states.is_contiguous() and hidden_states.device == qkvz.weight.device == ba.weight.device ) @@ -7229,7 +7256,7 @@ def qwen_gdn_input_projection_core( self.in_proj_ba.weight, ) _log_runtime_route_once( - "SM70 GDN QPN8 K5120/N4096 plus FP16 b/a N24 split route enabled." + "SM70 GDN QPN8 N4096 plus FP16 b/a N24 split route enabled." ) z = z_out else: @@ -7337,7 +7364,7 @@ def qwen_gdn_input_projection( self.in_proj_ba.weight, ) _log_runtime_route_once( - "SM70 GDN QPN8 K5120/N4096 plus FP16 b/a N24 split route enabled." + "SM70 GDN QPN8 N4096 plus FP16 b/a N24 split route enabled." ) return diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 935aa700f1..9344943de6 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -17,6 +17,7 @@ STR_DTYPE_TO_TORCH_DTYPE, get_kv_cache_torch_dtype, ) +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum logger = init_logger(__name__) @@ -199,9 +200,10 @@ def short_conv_state_shape( tp_world_size: int, intermediate_size: int, conv_kernel: int, + num_spec: int = 0, ) -> tuple[tuple[int, int]]: conv_dim = divide(intermediate_size, tp_world_size) - conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1) + conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1 + num_spec) return (conv_state_shape,) @classmethod @@ -301,6 +303,10 @@ class MambaCopySpec: num_accepted_tokens: int - number of accepted tokens used to compute the copy offset. Range: 1 .. 1 + num_speculative_tokens (inclusive). """ +MambaStateCopyFuncs: TypeAlias = tuple[MambaStateCopyFunc, ...] +MambaStateCopyFuncsByType: TypeAlias = dict[ + MambaAttentionBackendEnum, MambaStateCopyFuncs +] def get_conv_copy_spec( diff --git a/vllm/model_executor/layers/mamba/ops/gather_initial_states.py b/vllm/model_executor/layers/mamba/ops/gather_initial_states.py new file mode 100644 index 0000000000..7bf5de6e72 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/gather_initial_states.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +@triton.jit +def _gather_initial_states_kernel( + state_ptr, + indices_ptr, + has_initial_state_ptr, + output_ptr, + stride_state_batch, + stride_indices, + stride_has_initial_state, + row_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +): + block_idx = tl.program_id(0) + batch_idx = tl.program_id(1) + offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < row_size + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + has_initial_state = tl.load( + has_initial_state_ptr + batch_idx * stride_has_initial_state + ).to(tl.int1) + state_idx = tl.load(indices_ptr + batch_idx * stride_indices).to(tl.int64) + state_idx = tl.where(has_initial_state, state_idx, 0) + values = tl.load( + state_ptr + state_idx * stride_state_batch + offsets, + mask=mask & has_initial_state, + other=0.0, + ) + tl.store(output_ptr + batch_idx * row_size + offsets, values, mask=mask) + + +def gather_initial_states( + state: torch.Tensor, + indices: torch.Tensor, + has_initial_state: torch.Tensor, +) -> torch.Tensor: + """Gather state rows and zero rows without an initialized state.""" + assert state.ndim >= 2 and state.is_cuda + assert indices.ndim == 1 and has_initial_state.ndim == 1 + assert indices.shape == has_initial_state.shape + assert indices.device == state.device == has_initial_state.device + assert indices.dtype in (torch.int32, torch.int64) + assert has_initial_state.dtype == torch.bool + + row_size = state[0].numel() + assert state[0].is_contiguous() + output = torch.empty( + (indices.numel(), *state.shape[1:]), + dtype=state.dtype, + device=state.device, + ) + block_size = min(triton.next_power_of_2(row_size), 1024) + grid = (triton.cdiv(row_size, block_size), indices.numel()) + _gather_initial_states_kernel[grid]( + state, + indices, + has_initial_state, + output, + state.stride(0), + indices.stride(0), + has_initial_state.stride(0), + row_size=row_size, + BLOCK_SIZE=block_size, + num_warps=8, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return output diff --git a/vllm/model_executor/layers/mamba/ops/scatter_states.py b/vllm/model_executor/layers/mamba/ops/scatter_states.py new file mode 100644 index 0000000000..87340dbdb7 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/scatter_states.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +@triton.jit +def _scatter_states_kernel( + state_ptr, + src_ptr, + indices_ptr, + stride_state_batch, + stride_src_batch, + stride_indices, + row_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +): + block_idx = tl.program_id(0) + batch_idx = tl.program_id(1) + offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < row_size + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + state_idx = tl.load(indices_ptr + batch_idx * stride_indices).to(tl.int64) + values = tl.load(src_ptr + batch_idx * stride_src_batch + offsets, mask=mask) + tl.store(state_ptr + state_idx * stride_state_batch + offsets, values, mask=mask) + + +def scatter_states( + state: torch.Tensor, + src: torch.Tensor, + indices: torch.Tensor, +) -> None: + """Scatter dense source rows into unique state-cache slots.""" + assert state.ndim >= 2 and state.is_cuda + assert src.ndim == state.ndim and indices.ndim == 1 + assert indices.device == state.device + assert src.shape[1:] == state.shape[1:] + assert src.shape[0] == indices.shape[0] + assert indices.dtype in (torch.int32, torch.int64) + + row_size = state[0].numel() + assert state[0].is_contiguous() and src[0].is_contiguous() + block_size = min(triton.next_power_of_2(row_size), 1024) + grid = (triton.cdiv(row_size, block_size), indices.numel()) + _scatter_states_kernel[grid]( + state, + src, + indices, + state.stride(0), + src.stride(0), + indices.stride(0), + row_size=row_size, + BLOCK_SIZE=block_size, + num_warps=8, + launch_pdl=current_platform.is_arch_support_pdl(), + ) diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index 5249481293..857af315f8 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -6,11 +6,24 @@ # import vllm.model_executor.kernels.mhc # noqa: F401 import vllm.model_executor.kernels.mhc as mhc_kernels from vllm.model_executor.custom_op import CustomOp +from vllm.platforms import current_platform from vllm.utils.import_utils import has_tilelang HAS_TILELANG = has_tilelang() +def _use_sm70_fp16_mhc_fallback(residual: torch.Tensor, *, prefill_only: bool) -> bool: + del prefill_only + if residual.dtype != torch.float16 or not current_platform.is_cuda(): + return False + capability = current_platform.get_device_capability() + if capability is None or capability.to_int() != 70: + return False + # GLM-5.3's exact mHC4/H4096 contract has dedicated SM70 kernels for + # standalone pre/post, small-M fused decode, and large-M fused prefill. + return residual.shape[-2:] != (4, 4096) + + # --8<-- [start:mhc_pre] @CustomOp.register("mhc_pre") class MHCPreOp(CustomOp): @@ -41,6 +54,21 @@ def forward_cuda( norm_weight: torch.Tensor | None = None, norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if _use_sm70_fp16_mhc_fallback(residual, prefill_only=False): + return mhc_kernels.mhc_pre_torch( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + norm_weight, + norm_eps, + ) return torch.ops.vllm.mhc_pre_tilelang( residual, fn, @@ -170,6 +198,8 @@ def forward_cuda( post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor, ) -> torch.Tensor: + if _use_sm70_fp16_mhc_fallback(residual, prefill_only=False): + return mhc_kernels.mhc_post_torch(x, residual, post_layer_mix, comb_res_mix) return torch.ops.vllm.mhc_post_tilelang( x, residual, post_layer_mix, comb_res_mix ) @@ -335,6 +365,25 @@ def forward_cuda( norm_weight: torch.Tensor | None = None, norm_eps: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if _use_sm70_fp16_mhc_fallback(residual, prefill_only=True): + return mhc_kernels.mhc_fused_post_pre_torch( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + tile_n, + norm_weight, + norm_eps, + ) return torch.ops.vllm.mhc_fused_post_pre_tilelang( x, residual, @@ -396,3 +445,13 @@ def forward_native(self, *args, **kwargs): raise NotImplementedError( "Native implementation of mhc_fused_post_pre is not available" ) + + +def hc_expand(x: torch.Tensor, n: int) -> torch.Tensor: + """Expand ``[tokens, hidden]`` into residual streams.""" + return x.unsqueeze(1).expand(-1, n, -1).contiguous() + + +def hc_contract(x: torch.Tensor, n: int) -> torch.Tensor: + """Contract residual streams by averaging them.""" + return x.mean(dim=1) diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index 856f6bb8a3..2341173b33 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -8,6 +8,7 @@ from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.attention import MLAAttention from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.models.common.ops import fused_q_kv_rmsnorm @dataclass @@ -16,7 +17,7 @@ class MLAModules: kv_a_layernorm: torch.nn.Module kv_b_proj: torch.nn.Module - rotary_emb: torch.nn.Module + rotary_emb: torch.nn.Module | None o_proj: torch.nn.Module fused_qkv_a_proj: torch.nn.Module | None kv_a_proj_with_mqa: torch.nn.Module | None @@ -65,6 +66,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", skip_topk: bool = False, + fuse_qkv_rmsnorm: bool = False, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -93,6 +95,7 @@ def __init__( # the topk_tokens buffer written by a previous layer in the same pass. # Refer: https://arxiv.org/abs/2603.12201 for more details. self.skip_topk = skip_topk + self.fuse_qkv_rmsnorm = fuse_qkv_rmsnorm if self.indexer is not None: assert hasattr(self.indexer, "topk_tokens") self.topk_tokens = self.indexer.topk_tokens @@ -141,8 +144,8 @@ def forward( [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1, ) - q_c = self.q_a_layernorm(q_c) - q = self.q_b_proj(q_c)[0] + if not self.fuse_qkv_rmsnorm: + q_c = self.q_a_layernorm(q_c) else: assert self.kv_a_proj_with_mqa is not None, ( "kv_a_proj_with_mqa is required when q_lora_rank is None" @@ -154,7 +157,21 @@ def forward( q = self.q_proj(hidden_states)[0] kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - kv_c_normed = self.kv_a_layernorm(kv_c) + if self.fuse_qkv_rmsnorm and q_c is not None: + assert self.q_a_layernorm is not None + q_c, kv_c_normed = fused_q_kv_rmsnorm( + q_c, + kv_c, + self.q_a_layernorm.weight.data, + self.kv_a_layernorm.weight.data, + self.q_a_layernorm.variance_epsilon, + ) + else: + kv_c_normed = self.kv_a_layernorm(kv_c) + + if self.q_lora_rank is not None: + assert q_c is not None and self.q_b_proj is not None + q = self.q_b_proj(q_c)[0] q = q.view(-1, self.num_heads, self.qk_head_dim) # Add head dim of 1 to k_pe diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index ce0ed71b49..9b481a0241 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -480,15 +480,11 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: group_size=self.quant_config.group_size, ) layer.qweight = torch.nn.Parameter( - torch.empty( - 0, dtype=torch.int32, device=layer.qweight.device - ), + torch.empty(0, dtype=torch.int32, device=layer.qweight.device), requires_grad=False, ) layer.qzeros = torch.nn.Parameter( - torch.empty( - 0, dtype=torch.int32, device=layer.scales.device - ), + torch.empty(0, dtype=torch.int32, device=layer.scales.device), requires_grad=False, ) layer.scales = torch.nn.Parameter( @@ -498,9 +494,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: requires_grad=False, ) layer.g_idx = torch.nn.Parameter( - torch.empty( - 0, dtype=torch.int32, device=layer.scales.device - ), + torch.empty(0, dtype=torch.int32, device=layer.scales.device), requires_grad=False, ) logger.info_once("SM70 GPTQ TurboMind dense path enabled.") diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index 2766740c06..860daeb8d6 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -610,9 +610,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: torch.empty(0, dtype=tm_scales.dtype, device=tm_weight.device), requires_grad=False, ) - logger.info_once( - "SM70 AWQ TurboMind dense path enabled under AWQ Marlin." - ) + logger.info_once("SM70 AWQ TurboMind dense path enabled under AWQ Marlin.") return # AWQ checkpoints use a non-standard packing order and pack qweight diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_nvfp4.py index 1fa4f8632e..5f1fc7e1b9 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_nvfp4.py @@ -97,8 +97,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight, envs.VLLM_SM70_NVFP4_TURBOMIND ): logger.info_once( - "SM70 compressed-tensors NVFP4 TurboMind W4A16 dense path " - "enabled." + "SM70 compressed-tensors NVFP4 TurboMind W4A16 dense path enabled." ) sm70_tm.prepare_nvfp4_linear(layer) layer.weight = Parameter( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py index 07e252758f..f9d6f0b9b5 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py @@ -108,15 +108,11 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) sm70_tm.prepare_mxfp4_linear(layer) layer.weight_packed = Parameter( - torch.empty( - 0, dtype=torch.uint8, device=layer.weight_packed.device - ), + torch.empty(0, dtype=torch.uint8, device=layer.weight_packed.device), requires_grad=False, ) layer.weight_scale = Parameter( - torch.empty( - 0, dtype=torch.uint8, device=layer.weight_scale.device - ), + torch.empty(0, dtype=torch.uint8, device=layer.weight_scale.device), requires_grad=False, ) return diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py index c47e313ba9..0ff88ba608 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py @@ -87,6 +87,7 @@ def _missing_sm70_nvfp4_qpn4_ops() -> list[str]: "nvfp4_qpn2_gated_sm70_out", "nvfp4_qpn2_dispatch_sm70_out", ) +_SM70_NVFP4_QPN2_PREFILL_REQUIRED_OPS = ("nvfp4_qpn4_prefill_sm70_out",) def _is_qpn2_layer(layer: torch.nn.Module) -> bool: @@ -111,6 +112,14 @@ def _missing_qpn2_ops() -> list[str]: ] +def _missing_qpn2_prefill_ops() -> list[str]: + return [ + name + for name in _SM70_NVFP4_QPN2_PREFILL_REQUIRED_OPS + if not hasattr(torch.ops._C, name) + ] + + def _explicit_nvfp4_emulation_requested() -> bool: if envs.VLLM_USE_NVFP4_CT_EMULATIONS or envs.VLLM_NVFP4_GEMM_BACKEND == "emulation": return True @@ -305,6 +314,25 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight.data, layer.weight_scale.data ) qpn2_global_scale = float(layer.weight_global_scale.item()) + qpn2_prefill_workspace = None + if envs.VLLM_SM70_NVFP4_QPN2_PREFILL: + missing_prefill_ops = _missing_qpn2_prefill_ops() + if missing_prefill_ops: + logger.warning_once( + "The requested SM70 NVFP4 QPN2-packed prefill " + "route is unavailable; retaining TurboMind for " + f"large M. Missing ops: {missing_prefill_ops}." + ) + else: + qpn2_prefill_workspace = sm70_tm.get_nvfp4_qpn4_dense_workspace( + layer.weight + ) + if qpn2_prefill_workspace is None: + logger.warning_once( + "Insufficient memory for the bounded SM70 " + "NVFP4 QPN2-packed prefill workspace; " + "retaining TurboMind for large M." + ) use_gated_silu = bool( envs.VLLM_SM70_NVFP4_DENSE_GATED_SILU and is_qpn4_gate and not use_qpn2 @@ -329,10 +357,28 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.sm70_nvfp4_qpn2_split_k = split_k layer.sm70_nvfp4_qpn2_nacc = nacc layer.sm70_nvfp4_qpn2_gated_silu = suffix == "gate_up_proj" + layer.sm70_nvfp4_qpn2_prefill_dense_weight_ptr = ( + 0 + if qpn2_prefill_workspace is None + else qpn2_prefill_workspace.data_ptr() + ) + if qpn2_prefill_workspace is not None: + # QPN2 and QPN4 use the same physical tile order, but + # expose checkpoint-native [N, K/2] and GEMM-native + # [K, N/2] shapes respectively. Keep zero-copy views so + # prefill does not retain a third multi-GB weight layout. + layer.sm70_nvfp4_qpn2_prefill_codes = qpn2_codes.view(k, n // 2) + layer.sm70_nvfp4_qpn2_prefill_scales = qpn2_scales.view(k // 16, n) logger.info_once( "SM70 NVFP4 QPN2 M<=8 route enabled for a compatible " "TP4 projection contract." ) + if qpn2_prefill_workspace is not None: + logger.info_once( + "SM70 NVFP4 QPN2-packed bounded FP16 prefill route " + "enabled for M>=%d.", + envs.VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M, + ) elif use_gated_silu: logger.info_once( "SM70 NVFP4 TurboMind gated-SiLU single-layout path enabled." @@ -403,6 +449,26 @@ def _apply_qpn2( ) if x_2d.shape[0] == 0: return out_2d.reshape(*x.shape[:-1], output_size) + prefill_dense_weight_ptr = int( + getattr(layer, "sm70_nvfp4_qpn2_prefill_dense_weight_ptr", 0) + ) + if ( + x_2d.shape[0] >= envs.VLLM_SM70_NVFP4_QPN2_PREFILL_MIN_M + and prefill_dense_weight_ptr + ): + sm70_ops.nvfp4_qpn4_prefill_sm70_out( + out_2d, + prefill_dense_weight_ptr, + x_2d, + layer.sm70_nvfp4_qpn2_prefill_codes, + layer.sm70_nvfp4_qpn2_prefill_scales, + float(layer.sm70_nvfp4_qpn2_global_scale), + True, + gated_silu, + ) + if bias is not None: + out_2d.add_(bias) + return out_2d.reshape(*x.shape[:-1], output_size) state = getattr(layer, sm70_tm.STATE_ATTR) split_k = int(layer.sm70_nvfp4_qpn2_split_k) nacc = int(layer.sm70_nvfp4_qpn2_nacc) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index e81534de47..f438c291da 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -132,25 +132,25 @@ "qkv_proj": (5120, 3584), } _SM70_FP8_QPN8_PP2_TP4_CONFIGS = { - # Accuracy-first real-weight schedules: (K, N, fused gated-SiLU) maps to - # (split-K, accumulator chains, prefetch codes). The shared-expert gate/up - # projection is deliberately absent after its four-pattern comparison - # matched only 53.5--60.9% of TurboMind FP16 elements. - (4096, 1536, False): (16, 1, True), - (1024, 8192, False): (4, 1, False), - (2048, 4096, False): (8, 2, False), - (4096, 1024, False): (16, 1, True), - (512, 4096, False): (4, 1, False), + # Real-weight speed winners remain numerically bounded for every admitted + # projection. Shared-expert gate/up may use K4096/N1024 only through its + # separately gated non-fused contract; its fused activation remains unsafe. + (4096, 1536, False): (32, 2, False), + (1024, 8192, False): (8, 2, False), + (2048, 4096, False): (16, 2, False), + (4096, 1024, False): (32, 2, False), + (512, 4096, False): (16, 2, False), } _SM70_FP8_QPN8_PP2_TP4_SHAPES = { - # Operator role: accepted (layer TP size, K, N) tuples. The shared-expert - # gate/up projection is excluded by numerical evidence. The replicated - # indexer wq_b is excluded by TP size: its long-prefill work may overlap - # the main wq_b and cannot share one dense fallback workspace. + # Operator role: accepted (layer TP size, K, N) tuples. Gate/up has an + # additional exact shared-expert and explicit opt-in check below. The + # replicated indexer wq_b is excluded by TP size: its long-prefill work + # may overlap the main wq_b and cannot share one dense fallback workspace. "fused_wqa_wkv": {(1, 4096, 1536)}, "wq_b": {(4, 1024, 8192)}, "wo_b": {(4, 2048, 4096)}, "down_proj": {(4, 512, 4096)}, + "gate_up_proj": {(4, 4096, 1024)}, } _SM70_FP8_QPN8_PP2_TP4_WORKSPACE_ELEMENTS = max( k * n for k, n, _ in _SM70_FP8_QPN8_PP2_TP4_CONFIGS @@ -169,6 +169,61 @@ _sm70_fp8_qpn8_pp2_tp4_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} +def _is_sm70_fp8_pp2_tp4_shared_gate_layer(layer: torch.nn.Module) -> bool: + """Match the exact PP2 x TP4 shared-expert gate/up tensor.""" + prefix = str(getattr(layer, "prefix", "")) + return bool( + prefix.endswith(".shared_experts.gate_up_proj") + and int(getattr(layer, "tp_size", 1)) == 4 + and getattr(layer, "weight_block_size", None) == [128, 128] + and int(getattr(layer, "input_size_per_partition", 0)) == 4096 + and int(getattr(layer, "output_size_per_partition", 0)) == 1024 + and getattr(layer, "output_partition_sizes", None) == [512, 512] + and tuple(layer.weight.shape) == (1024, 4096) + ) + + +def _is_sm70_fp8_prescaled_m1_decode_layer(layer: torch.nn.Module) -> bool: + """Admit only the measured replicated fused WQA/WKV tensor contract.""" + return bool( + getattr(layer, "prefix", "").rsplit(".", 1)[-1] == "fused_wqa_wkv" + and int(getattr(layer, "tp_size", 1)) == 1 + and getattr(layer, "weight_block_size", None) == [128, 128] + and int(getattr(layer, "input_size_per_partition", 0)) == 4096 + and int(getattr(layer, "output_size_per_partition", 0)) == 1536 + and tuple(layer.weight.shape) == (1536, 4096) + ) + + +def _is_sm70_fp8_prescaled_m1_decode_runtime_contract() -> bool: + """Require the measured PP2 x TP4 no-spec single-request decode lane.""" + vllm_config = get_current_vllm_config() + parallel_config = vllm_config.parallel_config + scheduler_config = vllm_config.scheduler_config + return bool( + parallel_config.pipeline_parallel_size == 2 + and parallel_config.tensor_parallel_size == 4 + and scheduler_config.max_num_seqs == 1 + and not getattr(parallel_config, "enable_dbo", False) + and int(getattr(parallel_config, "ubatch_size", 0)) <= 1 + and getattr(vllm_config, "speculative_config", None) is None + ) + + +def _try_sm70_fp8_prescaled_decode_scales( + scales: torch.Tensor, +) -> torch.Tensor | None: + """Prescale only when the FP16 exponent shift is finite and reversible.""" + if scales.dtype != torch.float16 or not bool(torch.all(scales >= 0).item()): + return None + prescaled = scales.mul(256) + if not bool(torch.isfinite(prescaled).all().item()): + return None + if not torch.equal(prescaled.mul(1.0 / 256.0), scales): + return None + return prescaled + + def _is_sm70_fp8_exact_8k_prefill_layer(layer: torch.nn.Module) -> bool: if getattr(layer, "tp_size", 1) != 4: return False @@ -259,6 +314,13 @@ def _is_sm70_fp8_qpn8_pp2_tp4_runtime_contract() -> bool: ) +def _is_sm70_fp8_qpn8_pp2_tp4_shared_gate_contract( + layer: torch.nn.Module, +) -> bool: + """Match only the measured non-fused shared-expert gate/up tensor.""" + return _is_sm70_fp8_pp2_tp4_shared_gate_layer(layer) + + def _sm70_fp8_qpn8_pp2_tp4_config( layer: torch.nn.Module, *, gated_silu: bool ) -> tuple[int, int, bool] | None: @@ -266,6 +328,12 @@ def _sm70_fp8_qpn8_pp2_tp4_config( if getattr(layer, "weight_block_size", None) != [128, 128]: return None suffix = getattr(layer, "prefix", "").rsplit(".", 1)[-1] + if suffix == "gate_up_proj" and ( + gated_silu + or not envs.VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE + or not _is_sm70_fp8_qpn8_pp2_tp4_shared_gate_contract(layer) + ): + return None accepted = _SM70_FP8_QPN8_PP2_TP4_SHAPES.get(suffix) if accepted is None: return None @@ -839,7 +907,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: ) if missing_ops: logger.warning_once( - "The default SM70 PP2 x TP4 QPN8 route is unavailable " + "The requested SM70 PP2 x TP4 QPN8 route is unavailable " "in the loaded vllm._C; retaining TurboMind FP8." ) workspace = ( @@ -987,6 +1055,13 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: ) or (generic_qpn8_candidate and _is_sm70_fp8_qpn8_runtime_contract()) ) + nonfused_shared_gate = bool( + pp2_tp4_qpn8_candidate + and qpn8_runtime + and _is_sm70_fp8_qpn8_pp2_tp4_shared_gate_contract(layer) + ) + if nonfused_shared_gate: + use_gated_silu = False if qpn8_candidate_layer and not qpn8_runtime: logger.info_once( "The SM70 FP8 QPN8 route retains TurboMind unless its " @@ -1008,6 +1083,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: explicitly_enabled = ( os.getenv("VLLM_SM70_FP8_QPN8") == "1" or os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4") == "1" + or os.getenv("VLLM_SM70_FP8_QPN8_PP2_TP4_SHARED_GATE") == "1" ) if explicitly_enabled: raise RuntimeError( @@ -1015,7 +1091,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: f"source-built operators; missing: {missing_ops}." ) logger.warning_once( - "The default SM70 FP8 QPN8 route is unavailable in " + "The requested SM70 FP8 QPN8 route is unavailable in " "the loaded vllm._C; retaining the TurboMind layout." ) @@ -1068,6 +1144,11 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: "Default SM70 QPN8 enabled for the validated " "serialized PP2 x TP4 operator contract." ) + if nonfused_shared_gate: + logger.info_once( + "Default SM70 QPN8 shared-expert gate/up route " + "enabled with its external activation retained." + ) else: logger.info_once( "Memory-neutral SM70 FP8 QPN8 path enabled for " @@ -1080,6 +1161,29 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: "workspace; retaining the TurboMind layout." ) + prescaled_decode_requested = envs.VLLM_SM70_FP8_PRESCALED_M1_DECODE + prescaled_decode_explicit = ( + os.getenv("VLLM_SM70_FP8_PRESCALED_M1_DECODE") == "1" + ) + prescaled_decode_layer = _is_sm70_fp8_prescaled_m1_decode_layer(layer) + prescaled_decode_runtime = bool( + prescaled_decode_layer + and _is_sm70_fp8_prescaled_m1_decode_runtime_contract() + ) + if ( + prescaled_decode_explicit + and prescaled_decode_layer + and not (self.is_scale_e8m0 and prescaled_decode_runtime) + ): + raise RuntimeError( + "Explicit SM70 FP8 prescaled decode requires UE8M0 128x128 " + "scales and the PP2 x TP4 no-spec single-request contract." + ) + use_prescaled_m1_decode = bool( + prescaled_decode_requested + and self.is_scale_e8m0 + and prescaled_decode_runtime + ) tm_weight, tm_scales, meta = sm70_ops.fp8_sm70_prepare( weight, weight_scale_inv, @@ -1107,6 +1211,36 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: layer.register_buffer("sm70_fp8_meta", meta, persistent=False) layer.sm70_fp8_k_ld = int(meta[0].item()) layer.sm70_fp8_q_ld = int(meta[1].item()) + if use_prescaled_m1_decode: + has_prescaled_op = hasattr( + torch.ops._C, "fp8_gemm_sm70_prescaled_m1_out" + ) + prescaled_scales = ( + _try_sm70_fp8_prescaled_decode_scales(tm_scales) + if has_prescaled_op + else None + ) + if prescaled_scales is None and prescaled_decode_explicit: + raise RuntimeError( + "Explicit SM70 FP8 prescaled decode requires the " + "source-built operator and finite, reversible FP16 " + "scales after the 256x exponent shift." + ) + if prescaled_scales is None: + logger.warning_once( + "SM70 FP8 prescaled decode is unavailable for the " + "loaded extension or scale range; retaining the " + "ordinary TurboMind transform." + ) + else: + layer.register_buffer( + "sm70_fp8_decode_prescaled_scales", + prescaled_scales, + persistent=False, + ) + logger.info_once( + "Exact SM70 fused-WQA/WKV prescaled decode path enabled." + ) if ( envs.VLLM_SM70_FP8_PREFILL_FAST_SELECTOR and envs.VLLM_SM70_FP8_PREFILL_PRESCALED @@ -1392,6 +1526,9 @@ def apply( prefill_prescaled_scales = getattr( layer, "sm70_fp8_prefill_prescaled_scales", None ) + decode_prescaled_scales = getattr( + layer, "sm70_fp8_decode_prescaled_scales", None + ) prefill_min_m = getattr( layer, "sm70_fp8_prefill_exact_dense_min_m", @@ -1405,7 +1542,21 @@ def apply( gated_silu=False, min_prefill_m=prefill_min_m, ) - if visible_dense_out is not None: + if ( + decode_prescaled_scales is not None + and envs.VLLM_SM70_FP8_PRESCALED_M1_DECODE + and x_2d.shape[0] == 1 + ): + sm70_ops.fp8_gemm_sm70_prescaled_m1_out( + out_2d, + x_2d, + layer.weight, + decode_prescaled_scales, + 128, + layer.sm70_fp8_k_ld, + layer.sm70_fp8_q_ld, + ) + elif visible_dense_out is not None: out_2d = visible_dense_out elif prefill_workspace_ptr is not None and x_2d.dtype == torch.float16: sm70_ops.fp8_gemm_sm70_prefill_dispatch_out( diff --git a/vllm/model_executor/layers/quantization/input_quant_fp8.py b/vllm/model_executor/layers/quantization/input_quant_fp8.py index c32904339a..a6cbe801e6 100644 --- a/vllm/model_executor/layers/quantization/input_quant_fp8.py +++ b/vllm/model_executor/layers/quantization/input_quant_fp8.py @@ -227,9 +227,7 @@ def forward_native( else: x_max = x.abs().max().unsqueeze(-1).to(torch.float32) - scale = (x_max / self.fp8_max).clamp( - min=self.fp8_min_scaling_factor - ) + scale = (x_max / self.fp8_max).clamp(min=self.fp8_min_scaling_factor) else: scale = prep_scale_for_group_broadcast(scale, x, self.group_shape) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 09b3049800..35384dbdf0 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1161,6 +1161,28 @@ def get_name(self) -> QuantizationMethods: def get_supported_act_dtypes(self) -> list[torch.dtype]: return [torch.bfloat16, torch.half, torch.float8_e4m3fn] + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> "QuantizeMethodBase | None": + if ( + isinstance(layer, RoutedExperts) + and not self.is_layer_excluded(prefix) + and sm70_tm.is_exact_sm70_cuda_platform() + ): + if not sm70_tm.should_use_nvfp4_moe_turbomind(): + raise NotImplementedError( + "ModelOpt NVFP4 MoE on SM70 requires the TurboMind backend." + ) + from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( + ModelOptNvFp4SM70MoEMethod, + ) + + return ModelOptNvFp4SM70MoEMethod( + quant_config=self, + moe_config=layer.moe_config, + ) + return super().get_quant_method(layer, prefix) + @classmethod def get_min_capability(cls) -> int: # Do not unconditionally lower the class-wide gate to 70. W4A4 @@ -2548,6 +2570,19 @@ def get_quant_method( moe_config=layer.moe_config, ) if quant_algo == "NVFP4": + if sm70_tm.is_exact_sm70_cuda_platform(): + if not sm70_tm.should_use_nvfp4_moe_turbomind(): + raise NotImplementedError( + "ModelOpt NVFP4 MoE on SM70 requires the TurboMind backend." + ) + from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( + ModelOptNvFp4SM70MoEMethod, + ) + + return ModelOptNvFp4SM70MoEMethod( + quant_config=self.nvfp4_config, + moe_config=layer.moe_config, + ) return ModelOptNvFp4FusedMoE( quant_config=self.nvfp4_config, moe_config=layer.moe_config, diff --git a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py index fdd7bd202c..56e2cdf789 100644 --- a/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/mxfp4_sm70_moe.py @@ -9,6 +9,7 @@ from __future__ import annotations +import os from typing import Final import torch @@ -41,6 +42,30 @@ _DEEPSEEK_V4_FLASH_NUM_EXPERTS: Final = 256 _DEEPSEEK_V4_FLASH_TOP_K: Final = 6 _GRAPH_SAFE_MAX_TOKENS: Final = 8 +_MXFP4_QPN_M1_ENV: Final = "VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE" +_MXFP4_QPN_M1_OP: Final = "mxfp4_moe_qpn_m1_sm70_out" + + +def _mxfp4_qpn_m1_op_available() -> bool: + return hasattr(torch.ops._C, _MXFP4_QPN_M1_OP) + + +def _mxfp4_qpn_m1_extension_enabled() -> bool: + """Resolve the default-on route without breaking an older extension.""" + if not envs.VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE: + return False + if _mxfp4_qpn_m1_op_available(): + return True + if os.getenv(_MXFP4_QPN_M1_ENV) is not None: + raise RuntimeError( + "The explicitly enabled SM70 MXFP4 QPN M1 route requires the " + f"source-built {_MXFP4_QPN_M1_OP} operator." + ) + logger.warning_once( + "The default SM70 MXFP4 QPN M1 route is unavailable in the loaded " + "vllm._C; retaining the TurboMind dense-stage path." + ) + return False def _mxfp4_active_expert_b1_enabled() -> bool: @@ -188,15 +213,36 @@ def _select_mxfp4_stage_dispatch( def _select_mxfp4_direct_order_offsets( buffers: dict[str, torch.Tensor], ) -> torch.Tensor: - """Return the immutable one-row offsets for B1 direct-order decode. + """Return immutable one-row offsets for B1 direct-order decode. Active-expert compaction overwrites ``compact_expert_offsets`` during - M2--M8 warmup and verifier calls. Direct-order B1 has exactly one row per - route, so it must use the separately maintained slot offsets instead. + M2--M8 warmup and verifier calls. Direct-order B1 has one row per route, + so it must use the separately maintained slot offsets instead. """ return buffers["slot_expert_offsets"] +def _mxfp4_qpn_m1_decode_contract(layer: RoutedExperts, *, direct_order: bool) -> bool: + """Admit only the measured TP4 six-route W13/W2 tensor pair.""" + return bool( + envs.VLLM_SM70_MXFP4_MOE_QPN_M1_DECODE + and getattr(layer, "sm70_mxfp4_qpn_m1_available", False) + and direct_order + and int(layer.moe_config.tp_size) == 4 + and int(layer.local_num_experts) == 256 + and int(layer.global_num_experts) == 256 + and int(layer.sm70_mxfp4_w13_k_dim) == 4096 + and int(layer.sm70_mxfp4_w13_n_dim) == 1024 + and int(layer.sm70_mxfp4_w2_k_dim) == 512 + and int(layer.sm70_mxfp4_w2_n_dim) == 4096 + and int(layer.sm70_mxfp4_group_size) == 32 + and tuple(layer.w13_tm_weight.shape) == (256, 4096, 128) + and tuple(layer.w13_tm_scales.shape) == (256, 128, 1024) + and tuple(layer.w2_tm_weight.shape) == (256, 512, 512) + and tuple(layer.w2_tm_scales.shape) == (256, 16, 4096) + ) + + def validate_mxfp4_sm70_moe_contract( *, global_num_experts: int, @@ -366,6 +412,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: "DeepSeek-V4 MXFP4 MoE on SM70 requires the TurboMind CUDA " "extension with " + ", ".join(missing_ops) + "." ) + layer.sm70_mxfp4_qpn_m1_available = _mxfp4_qpn_m1_extension_enabled() if not hasattr(torch.ops._moe_C, "moe_permute_with_scratch"): raise RuntimeError( "DeepSeek-V4 MXFP4 MoE graph-safe B1 requires " @@ -442,6 +489,21 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: w13_q_ld = int(w13_meta[0][1].item()) w2_k_ld = int(w2_meta[0][0].item()) w2_q_ld = int(w2_meta[0][1].item()) + + # Stacking 256 experts temporarily holds both the per-expert tensors + # and their contiguous replacements. Release those load-only tensors + # before the pointer tables allocate so V100 does not exhaust driver + # memory while the CUDA caching allocator retains the stack workspace. + del w13_tm_weights, w13_tm_scales, w13_meta + del w2_tm_weights, w2_tm_scales, w2_meta + del w13_packed, w13_scales, prepared_w13 + del w2_packed, w2_scales, prepared_w2 + del layer.w13_weight + del layer.w13_weight_scale + del layer.w2_weight + del layer.w2_weight_scale + torch.accelerator.empty_cache() + w13_ptrs = sm70_ops.awq_moe_build_strided_ptrs( layer.w13_tm_weight, layer.w13_tm_scales, @@ -472,12 +534,6 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: layer.sm70_mxfp4_group_size = MXFP4_GROUP_SIZE self._allocate_graph_safe_decode_buffers(layer) - # Raw checkpoint tensors are replaced by the equivalent TurboMind - # packed e2m1/UE8M0 representation, never by dequantized weights. - del layer.w13_weight - del layer.w13_weight_scale - del layer.w2_weight - del layer.w2_weight_scale logger.info_once( "SM70 TurboMind MXFP4 MoE enabled for DeepSeek-V4-Flash " "(local_experts=%d, graph_safe_decode=B1-B%d, " @@ -749,6 +805,38 @@ def apply( and topk_ids.dtype == torch.int32 and topk_ids.is_contiguous() ) + if _mxfp4_qpn_m1_decode_contract(layer, direct_order=direct_order): + route_ids = topk_ids.view(-1) + sm70_ops.mxfp4_moe_qpn_m1_sm70_out( + buffers["gate_up"], + x, + layer.w13_tm_weight, + layer.w13_tm_scales, + route_ids, + True, + ) + self._apply_swiglu(layer, buffers["intermediate"], buffers["gate_up"]) + sm70_ops.mxfp4_moe_qpn_m1_sm70_out( + buffers["sorted_output"], + buffers["intermediate"], + layer.w2_tm_weight, + layer.w2_tm_scales, + route_ids, + False, + ) + sm70_ops.awq_moe_single_token_weighted_reduce_out( + buffers["sorted_output"], + topk_weights, + buffers["token_expert_indices"], + output, + _DEEPSEEK_V4_FLASH_TOP_K, + layer.sm70_mxfp4_hidden_size, + ) + logger.info_once( + "Default SM70 MXFP4 QPN M1 route enabled for the exact " + "TP4 six-route W13/W2 tensor contract." + ) + return output if direct_order: route_ids = topk_ids.view(-1) route_offsets = _select_mxfp4_direct_order_offsets(buffers) diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index 7eb31d9ce3..5324dcb6ad 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Native SM70 TurboMind NVFP4 MoE for Qwen3.6-35B-A3B. +"""Native SM70 TurboMind NVFP4 MoE for validated expert shapes. The route keeps ModelOpt W4A16_NVFP4 expert weights packed. It combines the checkpoint's FP8 block scales with its explicit ModelOpt global scales once at @@ -16,6 +16,7 @@ from torch.nn import Parameter from vllm import _sm70_ops as sm70_ops +from vllm import envs from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( FusedMoEConfig, @@ -38,13 +39,126 @@ logger = init_logger(__name__) -_QWEN36_HIDDEN_SIZE: Final = 2048 -_QWEN36_INTERMEDIATE_SIZE: Final = 512 -_QWEN36_NUM_EXPERTS: Final = 256 -_QWEN36_TOP_K: Final = 8 -_QWEN36_SUPPORTED_TP_SIZES: Final = (1, 2, 4) +_SUPPORTED_CONTRACTS: Final = { + # (hidden size, global expert intermediate size, experts, top-k) + (2048, 512, 256, 8), # Qwen3.6-35B-A3B + (2560, 640, 512, 10), # Qwen3.8-Flash-Next + (4096, 2048, 288, 8), # GLM-5.3-Flash +} +_SUPPORTED_TP_SIZES: Final = (1, 2, 4) _GRAPH_SAFE_MAX_TOKENS: Final = 18 _COMPACT_GROUPED_MAX_TOKENS: Final = 10 +_MAX_SUPPORTED_TOP_K: Final = max(contract[3] for contract in _SUPPORTED_CONTRACTS) +_QWEN38_QPN_M1_W13_SPLIT_K: Final = 8 +_QWEN38_QPN_M1_W2_SPLIT_K: Final = 1 + + +def _use_qwen38_qpn_m1_decode( + layer: RoutedExperts, + x: torch.Tensor, + topk_ids: torch.Tensor, +) -> bool: + """Admit only the exact validated Qwen3.8 TP4 single-token route.""" + return bool( + envs.VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE + and x.shape == (1, 2560) + and x.dtype == torch.float16 + and x.is_contiguous() + and topk_ids.shape == (1, 10) + and topk_ids.dtype == torch.int32 + and topk_ids.is_contiguous() + and int(layer.moe_config.tp_size) == 4 + and int(layer.sm70_nvfp4_num_experts) == 512 + and int(layer.sm70_nvfp4_hidden_size) == 2560 + and int(layer.sm70_nvfp4_intermediate_size) == 160 + and int(layer.sm70_nvfp4_top_k) == 10 + ) + + +@triton.jit +def _prepare_single_token_slots_kernel( + input_ptr, + topk_ids_ptr, + expanded_input_ptr, + active_expert_ids_ptr, + HIDDEN: tl.constexpr, + BLOCK: tl.constexpr, +): + slot = tl.program_id(0) + offsets = tl.arange(0, BLOCK) + mask = offsets < HIDDEN + values = tl.load(input_ptr + offsets, mask=mask, other=0.0) + tl.store(expanded_input_ptr + slot * HIDDEN + offsets, values, mask=mask) + expert_id = tl.load(topk_ids_ptr + slot) + tl.store(active_expert_ids_ptr + slot, expert_id.to(tl.int32)) + + +def _prepare_single_token_slots( + x: torch.Tensor, + topk_ids: torch.Tensor, + expanded_input: torch.Tensor, + active_expert_ids: torch.Tensor, +) -> None: + top_k = topk_ids.numel() + hidden = x.shape[1] + if x.shape[0] != 1 or tuple(topk_ids.shape) != (1, top_k): + raise ValueError("SM70 NVFP4 direct routing requires one input token.") + if tuple(expanded_input.shape) != (top_k, hidden): + raise ValueError("SM70 NVFP4 direct routing buffer shape mismatch.") + if active_expert_ids.numel() != top_k: + raise ValueError("SM70 NVFP4 direct expert-ID buffer shape mismatch.") + _prepare_single_token_slots_kernel[(top_k,)]( + x, + topk_ids, + expanded_input, + active_expert_ids, + HIDDEN=hidden, + BLOCK=triton.next_power_of_2(hidden), + num_warps=8, + ) + + +@triton.jit +def _single_token_weighted_reduce_kernel( + expert_output_ptr, + topk_weights_ptr, + output_ptr, + HIDDEN: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < HIDDEN + acc = tl.zeros((BLOCK,), tl.float32) + for slot in tl.static_range(0, TOP_K): + values = tl.load( + expert_output_ptr + slot * HIDDEN + offsets, + mask=mask, + other=0.0, + ) + weight = tl.load(topk_weights_ptr + slot) + acc += values.to(tl.float32) * weight + tl.store(output_ptr + offsets, acc, mask=mask) + + +def _single_token_weighted_reduce( + expert_output: torch.Tensor, + topk_weights: torch.Tensor, + output: torch.Tensor, +) -> None: + top_k, hidden = expert_output.shape + if tuple(topk_weights.shape) != (1, top_k) or tuple(output.shape) != (1, hidden): + raise ValueError("SM70 NVFP4 direct weighted-reduce shape mismatch.") + block = 256 + _single_token_weighted_reduce_kernel[(triton.cdiv(hidden, block),)]( + expert_output, + topk_weights, + output, + HIDDEN=hidden, + TOP_K=top_k, + BLOCK=block, + num_warps=4, + ) @triton.jit @@ -80,7 +194,7 @@ def _prepare_compact_slot_groups( active_expert_ids: torch.Tensor, ) -> None: total_slots = sorted_expert_ids.numel() - max_slots = _COMPACT_GROUPED_MAX_TOKENS * _QWEN36_TOP_K + max_slots = _COMPACT_GROUPED_MAX_TOKENS * _MAX_SUPPORTED_TOP_K if not (0 < total_slots <= max_slots): raise ValueError(f"Unsupported SM70 NVFP4 active-expert slots: {total_slots}") block = triton.next_power_of_2(total_slots + 1) @@ -100,25 +214,10 @@ def _prepare_compact_slot_groups( def validate_nvfp4_sm70_moe_contract(moe: FusedMoEConfig) -> None: """Reject every topology outside the validated SM70 NVFP4 contract.""" - if moe.num_experts != _QWEN36_NUM_EXPERTS: - raise NotImplementedError( - "SM70 TurboMind NVFP4 MoE currently supports Qwen3.6-35B-A3B " - f"with {_QWEN36_NUM_EXPERTS} experts, got {moe.num_experts}." - ) - if moe.experts_per_token != _QWEN36_TOP_K: - raise NotImplementedError( - "SM70 TurboMind NVFP4 MoE currently supports top-k=" - f"{_QWEN36_TOP_K}, got {moe.experts_per_token}." - ) - if moe.hidden_dim != _QWEN36_HIDDEN_SIZE: - raise NotImplementedError( - "SM70 TurboMind NVFP4 MoE currently supports hidden size " - f"{_QWEN36_HIDDEN_SIZE}, got {moe.hidden_dim}." - ) - if moe.tp_size not in _QWEN36_SUPPORTED_TP_SIZES: + if moe.tp_size not in _SUPPORTED_TP_SIZES: raise NotImplementedError( "SM70 TurboMind NVFP4 MoE currently supports tensor parallel " - f"sizes {_QWEN36_SUPPORTED_TP_SIZES}, got {moe.tp_size}." + f"sizes {_SUPPORTED_TP_SIZES}, got {moe.tp_size}." ) local_intermediate = moe.intermediate_size_per_partition if local_intermediate <= 0 or local_intermediate % NVFP4_GROUP_SIZE: @@ -126,11 +225,19 @@ def validate_nvfp4_sm70_moe_contract(moe: FusedMoEConfig) -> None: "SM70 TurboMind NVFP4 MoE requires a positive local intermediate " f"size divisible by {NVFP4_GROUP_SIZE}, got {local_intermediate}." ) - if local_intermediate * max(moe.tp_size, 1) != _QWEN36_INTERMEDIATE_SIZE: + global_intermediate = local_intermediate * max(moe.tp_size, 1) + contract = ( + moe.hidden_dim, + global_intermediate, + moe.num_experts, + moe.experts_per_token, + ) + if contract not in _SUPPORTED_CONTRACTS: raise NotImplementedError( - "SM70 TurboMind NVFP4 MoE currently supports Qwen3.6 expert " - f"intermediate size {_QWEN36_INTERMEDIATE_SIZE}; got local=" - f"{local_intermediate}, tp_size={moe.tp_size}." + "SM70 TurboMind NVFP4 MoE shape is not validated: " + f"hidden={moe.hidden_dim}, intermediate={global_intermediate}, " + f"experts={moe.num_experts}, top_k={moe.experts_per_token}. " + f"Validated contracts: {sorted(_SUPPORTED_CONTRACTS)}." ) if moe.moe_parallel_config.use_all2all_kernels: raise NotImplementedError( @@ -175,7 +282,7 @@ def _validate_weight_layout(layer: RoutedExperts) -> None: class ModelOptNvFp4SM70MoEMethod(ModelOptNvFp4FusedMoE): - """Qwen3.6 ModelOpt W4A16_NVFP4 experts on native TurboMind SM70.""" + """ModelOpt NVFP4 experts with FP16 activations on native SM70.""" def __init__( self, @@ -183,10 +290,10 @@ def __init__( moe_config: FusedMoEConfig, ) -> None: FusedMoEMethodBase.__init__(self, moe_config) - if quant_config.quant_method != "W4A16_NVFP4": + if quant_config.quant_method not in {"NVFP4", "W4A16_NVFP4"}: raise NotImplementedError( - "SM70 TurboMind ModelOpt NVFP4 MoE currently requires " - "W4A16_NVFP4 checkpoint weights." + "SM70 TurboMind ModelOpt NVFP4 MoE requires NVFP4-family " + f"checkpoint weights, got {quant_config.quant_method}." ) self.quant_config = quant_config self.use_a16 = True @@ -213,15 +320,18 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: "awq_moe_build_strided_ptrs", ) missing = [name for name in required_ops if not hasattr(torch.ops._C, name)] + if ( + envs.VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE + and not sm70_ops.has_nvfp4_qpn_m1_dispatch() + ): + missing.append("nvfp4_moe_qpn_m1_sm70_out") if missing: raise RuntimeError( - "Qwen3.6 NVFP4 MoE on SM70 requires the TurboMind extension " + "SM70 NVFP4 MoE requires the TurboMind extension " "with " + ", ".join(missing) + "." ) if not hasattr(torch.ops._moe_C, "moe_permute_with_scratch"): - raise RuntimeError( - "Qwen3.6 NVFP4 MoE on SM70 requires graph-safe MoE permute ops." - ) + raise RuntimeError("SM70 NVFP4 MoE requires graph-safe MoE permute ops.") if self.moe.has_bias: raise NotImplementedError("SM70 NVFP4 MoE does not support expert bias.") if layer.activation != MoEActivation.SILU: @@ -316,6 +426,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: layer.sm70_nvfp4_num_experts = num_experts layer.sm70_nvfp4_hidden_size = hidden layer.sm70_nvfp4_intermediate_size = intermediate + layer.sm70_nvfp4_top_k = int(layer.moe_config.experts_per_token) layer.sm70_nvfp4_w13_k_dim = hidden layer.sm70_nvfp4_w13_n_dim = 2 * intermediate layer.sm70_nvfp4_w2_k_dim = intermediate @@ -334,17 +445,21 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: del layer.w2_weight_scale_2 del layer.w2_input_scale logger.info_once( - "SM70 ModelOpt NVFP4 TurboMind MoE path enabled for " - "Qwen3.6-35B-A3B (local_experts=%d, graph_safe_decode=B1-B%d, " - "compact_grouped_decode=B1-B%d).", + "SM70 ModelOpt NVFP4 TurboMind MoE path enabled " + "(hidden=%d, local_intermediate=%d, local_experts=%d, top_k=%d, " + "graph_safe_decode=B1-B%d, compact_grouped_decode=B1-B%d).", + hidden, + intermediate, num_experts, + layer.sm70_nvfp4_top_k, _GRAPH_SAFE_MAX_TOKENS, _COMPACT_GROUPED_MAX_TOKENS, ) def _allocate_graph_safe_decode_buffers(self, layer: RoutedExperts) -> None: device = layer.w13_tm_weight.device - max_slots = _GRAPH_SAFE_MAX_TOKENS * _QWEN36_TOP_K + top_k = int(layer.sm70_nvfp4_top_k) + max_slots = _GRAPH_SAFE_MAX_TOKENS * top_k experts = int(layer.sm70_nvfp4_num_experts) hidden = int(layer.sm70_nvfp4_hidden_size) intermediate = int(layer.sm70_nvfp4_intermediate_size) @@ -372,19 +487,19 @@ def _allocate_graph_safe_decode_buffers(self, layer: RoutedExperts) -> None: ) layer._nvfp4_sm70_inv_permuted_idx = torch.empty( _GRAPH_SAFE_MAX_TOKENS, - _QWEN36_TOP_K, + top_k, dtype=torch.int32, device=device, ) layer._nvfp4_sm70_topk_ids = torch.empty( _GRAPH_SAFE_MAX_TOKENS, - _QWEN36_TOP_K, + top_k, dtype=torch.int32, device=device, ) layer._nvfp4_sm70_token_expert_indices = torch.arange( max_slots, dtype=torch.int32, device=device - ).view(_GRAPH_SAFE_MAX_TOKENS, _QWEN36_TOP_K) + ).view(_GRAPH_SAFE_MAX_TOKENS, top_k) layer._nvfp4_sm70_permuted_idx = torch.empty( max_slots, dtype=torch.int32, device=device ) @@ -406,7 +521,7 @@ def _allocate_graph_safe_decode_buffers(self, layer: RoutedExperts) -> None: layer._nvfp4_sm70_dense_expert_ids = torch.arange( experts, dtype=torch.int32, device=device ) - layer._nvfp4_sm70_compact_offsets = torch.empty( + layer._nvfp4_sm70_compact_offsets = torch.arange( max_slots + 1, dtype=torch.int32, device=device ) layer._nvfp4_sm70_active_expert_ids = torch.empty( @@ -417,7 +532,7 @@ def _allocate_graph_safe_decode_buffers(self, layer: RoutedExperts) -> None: def _persistent_buffers( layer: RoutedExperts, num_tokens: int ) -> dict[str, torch.Tensor]: - slots = num_tokens * _QWEN36_TOP_K + slots = num_tokens * int(layer.sm70_nvfp4_top_k) return { "output": layer._nvfp4_sm70_output[:num_tokens], "permuted_input": layer._nvfp4_sm70_permuted_input[:slots], @@ -446,7 +561,8 @@ def _eager_buffers( layer: RoutedExperts, num_tokens: int ) -> dict[str, torch.Tensor]: device = layer.w13_tm_weight.device - slots = num_tokens * _QWEN36_TOP_K + top_k = int(layer.sm70_nvfp4_top_k) + slots = num_tokens * top_k experts = int(layer.sm70_nvfp4_num_experts) hidden = int(layer.sm70_nvfp4_hidden_size) intermediate = int(layer.sm70_nvfp4_intermediate_size) @@ -476,14 +592,14 @@ def _eager_buffers( experts + 1, dtype=torch.int64, device=device ), "inv_permuted_idx": torch.empty( - num_tokens, _QWEN36_TOP_K, dtype=torch.int32, device=device + num_tokens, top_k, dtype=torch.int32, device=device ), "topk_ids": torch.empty( - num_tokens, _QWEN36_TOP_K, dtype=torch.int32, device=device + num_tokens, top_k, dtype=torch.int32, device=device ), "token_expert_indices": torch.arange( slots, dtype=torch.int32, device=device - ).view(num_tokens, _QWEN36_TOP_K), + ).view(num_tokens, top_k), "permuted_idx": torch.empty(slots, dtype=torch.int32, device=device), "sort_workspace": torch.empty( workspace_size, dtype=torch.int8, device=device @@ -492,7 +608,9 @@ def _eager_buffers( "sorted_row_idx": torch.empty(slots, dtype=torch.int32, device=device), "topk_ids_for_sort": torch.empty(slots, dtype=torch.int32, device=device), "dense_expert_ids": layer._nvfp4_sm70_dense_expert_ids, - "compact_offsets": torch.empty(slots + 1, dtype=torch.int32, device=device), + "compact_offsets": torch.arange( + slots + 1, dtype=torch.int32, device=device + ), "active_expert_ids": torch.empty(slots, dtype=torch.int32, device=device), } @@ -528,13 +646,18 @@ def apply( raise TypeError("SM70 NVFP4 MoE requires CUDA FP16 activations [M, H].") if not is_exact_sm70_cuda(x, enabled=True): raise RuntimeError("SM70 NVFP4 MoE dispatch is restricted to CUDA SM70.") - if x.shape[1] != _QWEN36_HIDDEN_SIZE: + hidden = int(layer.sm70_nvfp4_hidden_size) + top_k = int(layer.sm70_nvfp4_top_k) + if x.shape[1] != hidden: raise ValueError( "SM70 NVFP4 MoE activation hidden size mismatch: expected " - f"{_QWEN36_HIDDEN_SIZE}, got {x.shape[1]}." + f"{hidden}, got {x.shape[1]}." + ) + if tuple(topk_ids.shape) != (x.shape[0], top_k): + raise ValueError( + "SM70 NVFP4 MoE top-k ID shape mismatch: expected " + f"{(x.shape[0], top_k)}, got {tuple(topk_ids.shape)}." ) - if tuple(topk_ids.shape) != (x.shape[0], _QWEN36_TOP_K): - raise ValueError("SM70 NVFP4 MoE requires top-k IDs with shape [M, 8].") if tuple(topk_weights.shape) != tuple(topk_ids.shape): raise ValueError("SM70 NVFP4 MoE top-k weights and IDs must share shape.") if topk_weights.dtype != torch.float32: @@ -542,34 +665,77 @@ def apply( num_tokens = x.shape[0] if num_tokens == 0: - return x.new_empty((0, _QWEN36_HIDDEN_SIZE)) + return x.new_empty((0, hidden)) buffers = self._get_buffers(layer, num_tokens) output = buffers["output"] - output.zero_() - slots = num_tokens * _QWEN36_TOP_K - topk_ids_i32 = buffers["topk_ids"] - topk_ids_i32.copy_(topk_ids, non_blocking=True) - buffers["permuted_idx"].fill_(slots) - torch.ops._moe_C.moe_permute_with_scratch( - x, - topk_ids_i32, - buffers["token_expert_indices"], - layer.expert_map, - layer.global_num_experts, - layer.local_num_experts, - _QWEN36_TOP_K, - buffers["permuted_input"], - buffers["expert_offsets64"], - buffers["inv_permuted_idx"], - buffers["permuted_idx"], - buffers["sort_workspace"], - buffers["permuted_experts_id"], - buffers["sorted_row_idx"], - buffers["topk_ids_for_sort"], - ) - buffers["expert_offsets"].copy_(buffers["expert_offsets64"], non_blocking=True) + slots = num_tokens * top_k + direct_single_token = num_tokens == 1 + if _use_qwen38_qpn_m1_decode(layer, x, topk_ids): + logger.info_once( + "SM70 Qwen3.8 NVFP4 direct QPN-M1 expert path enabled " + "(TP4, E512/K10, W13 split8, W2 split1)." + ) + route_ids = topk_ids.view(-1) + sm70_ops.nvfp4_moe_qpn_m1_sm70_out( + buffers["gate_up"], + x, + layer.w13_tm_weight, + layer.w13_tm_scales, + route_ids, + True, + _QWEN38_QPN_M1_W13_SPLIT_K, + ) + self._apply_swiglu(layer, buffers["intermediate"], buffers["gate_up"]) + sm70_ops.nvfp4_moe_qpn_m1_sm70_out( + buffers["sorted_output"], + buffers["intermediate"], + layer.w2_tm_weight, + layer.w2_tm_scales, + route_ids, + False, + _QWEN38_QPN_M1_W2_SPLIT_K, + ) + _single_token_weighted_reduce( + buffers["sorted_output"], topk_weights, output + ) + return output + if direct_single_token: + _prepare_single_token_slots( + x, + topk_ids, + buffers["permuted_input"], + buffers["active_expert_ids"], + ) + stage_offsets = buffers["compact_offsets"] + stage_expert_ids = buffers["active_expert_ids"] + stage_experts = top_k + else: + output.zero_() + topk_ids_i32 = buffers["topk_ids"] + topk_ids_i32.copy_(topk_ids, non_blocking=True) + buffers["permuted_idx"].fill_(slots) + torch.ops._moe_C.moe_permute_with_scratch( + x, + topk_ids_i32, + buffers["token_expert_indices"], + layer.expert_map, + layer.global_num_experts, + layer.local_num_experts, + top_k, + buffers["permuted_input"], + buffers["expert_offsets64"], + buffers["inv_permuted_idx"], + buffers["permuted_idx"], + buffers["sort_workspace"], + buffers["permuted_experts_id"], + buffers["sorted_row_idx"], + buffers["topk_ids_for_sort"], + ) + buffers["expert_offsets"].copy_( + buffers["expert_offsets64"], non_blocking=True + ) - if num_tokens <= _COMPACT_GROUPED_MAX_TOKENS: + if not direct_single_token and num_tokens <= _COMPACT_GROUPED_MAX_TOKENS: _prepare_compact_slot_groups( buffers["permuted_experts_id"], buffers["compact_offsets"], @@ -578,7 +744,7 @@ def apply( stage_offsets = buffers["compact_offsets"] stage_expert_ids = buffers["active_expert_ids"] stage_experts = slots - else: + elif not direct_single_token: stage_offsets = buffers["expert_offsets"] stage_expert_ids = buffers["dense_expert_ids"] stage_experts = int(layer.sm70_nvfp4_num_experts) @@ -608,14 +774,19 @@ def apply( layer.sm70_nvfp4_w2_n_dim, layer.sm70_nvfp4_group_size, ) - torch.ops._moe_C.moe_unpermute( - buffers["sorted_output"], - topk_weights, - buffers["inv_permuted_idx"], - buffers["expert_offsets64"], - _QWEN36_TOP_K, - output, - ) + if direct_single_token: + _single_token_weighted_reduce( + buffers["sorted_output"], topk_weights, output + ) + else: + torch.ops._moe_C.moe_unpermute( + buffers["sorted_output"], + topk_weights, + buffers["inv_permuted_idx"], + buffers["expert_offsets64"], + top_k, + output, + ) return output def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/sm70_online_qpn8.py b/vllm/model_executor/layers/quantization/sm70_online_qpn8.py new file mode 100644 index 0000000000..6ab67cc784 --- /dev/null +++ b/vllm/model_executor/layers/quantization/sm70_online_qpn8.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Narrow online QPN8 route for Qwen4Exp FP16 decode projections on SM70.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +import vllm.envs as envs +from vllm import _sm70_ops as sm70_ops +from vllm.config import get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.logger import init_logger +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from torch import nn + + +logger = init_logger(__name__) + +_STATE_ATTR = "_sm70_qwen4_exp_online_qpn8" +_CODES_ATTR = "_sm70_qwen4_exp_online_qpn8_codes" +_SCALES_ATTR = "_sm70_qwen4_exp_online_qpn8_scales" +_HC_PARTIALS_ATTR = "_sm70_qwen4_exp_online_qpn8_hc_partials" +_WORKSPACE_ELEMENTS = 4096 * 2560 +_workspaces: dict[tuple[int, torch.dtype], torch.Tensor] = {} +_hc_partials: dict[int, torch.Tensor] = {} + +_HC_DOWN_SUFFIX = ".input_mix_weight_down_block_inject" +_HC_DOWN_LOGICAL_N = 336 +_HC_DOWN_PADDED_N = 384 + +_REQUIRED_OPS = ( + "fp8_qpn8_prepare_sm70", + "fp8_qpn8_dispatch_sm70_out", +) + + +def _shape_config(prefix: str, k: int, n: int) -> tuple[int, int, bool] | None: + """Return the measured ``split_k, nacc, prefetch`` for one local weight.""" + if prefix.endswith(_HC_DOWN_SUFFIX) and (k, n) == ( + 10240, + _HC_DOWN_LOGICAL_N, + ): + return 32, 1, False + if prefix.endswith(".input_mix_weight_up") and (k, n) == (320, 10240): + return 4, 2, False + if prefix.endswith(".linear_attn.in_proj_qkvz") and (k, n) == (2560, 4096): + return 16, 1, False + if prefix.endswith(".linear_attn.out_proj") and (k, n) == (1536, 2560): + return 12, 2, False + if prefix.endswith(".self_attn.qkv_proj") and (k, n) == (2560, 3584): + return 16, 1, False + if prefix.endswith(".self_attn.o_proj") and (k, n) == (1536, 2560): + return 12, 2, False + return None + + +def _exact_runtime_contract() -> bool: + try: + config = get_current_vllm_config() + text_config = config.model_config.hf_text_config + tp_size = get_tensor_model_parallel_world_size() + except (AssertionError, AttributeError, RuntimeError): + return False + return bool( + tp_size == 4 + and config.speculative_config is None + and int(getattr(text_config, "hidden_size", 0)) == 2560 + and int(getattr(text_config, "num_hidden_layers", 0)) == 48 + and int(getattr(text_config, "num_experts", 0)) == 512 + and int(getattr(text_config, "num_experts_per_tok", 0)) == 10 + and int(getattr(text_config, "moe_intermediate_size", 0)) == 640 + and int(getattr(text_config, "hc_count", 0)) == 4 + and int(getattr(text_config, "hc_lowrank", 0)) == 320 + and int(getattr(text_config, "num_attention_heads", 0)) == 24 + and int(getattr(text_config, "num_key_value_heads", 0)) == 2 + and int(getattr(text_config, "indexer_head_dim", 0)) == 128 + and int(getattr(text_config, "indexer_budget", 0)) == 2048 + and int(getattr(text_config, "indexer_compress_ratio", 0)) == 4 + ) + + +def _workspace(weight: torch.Tensor) -> torch.Tensor: + device_index = weight.device.index + if device_index is None: + device_index = torch.accelerator.current_device_index() + key = (device_index, weight.dtype) + workspace = _workspaces.get(key) + if workspace is None: + workspace = torch.empty( + (_WORKSPACE_ELEMENTS,), dtype=torch.float16, device=weight.device + ) + _workspaces[key] = workspace + return workspace + + +def _hc_workspace(weight: torch.Tensor) -> torch.Tensor: + device_index = weight.device.index + if device_index is None: + device_index = torch.accelerator.current_device_index() + partials = _hc_partials.get(device_index) + if partials is None: + partials = torch.empty( + (32 * _HC_DOWN_PADDED_N,), + dtype=torch.float32, + device=weight.device, + ) + _hc_partials[device_index] = partials + return partials + + +def maybe_prepare_online_qpn8(layer: nn.Module) -> bool: + """Quantize one admitted FP16 weight and retain only its QPN8 layout.""" + if getattr(layer, _STATE_ATTR, False): + return True + if not envs.VLLM_SM70_QWEN4_EXP_ONLINE_QPN8 or not _exact_runtime_contract(): + return False + if any(not hasattr(torch.ops._C, name) for name in _REQUIRED_OPS): + raise RuntimeError( + "VLLM_SM70_QWEN4_EXP_ONLINE_QPN8=1 requires the SM70 QPN8 ops" + ) + weight = getattr(layer, "weight", None) + if ( + weight is None + or weight.dtype != torch.float16 + or not weight.is_cuda + or weight.ndim != 2 + or not current_platform.is_device_capability(70) + ): + return False + + logical_n, k = (int(dim) for dim in weight.shape) + prefix = str(getattr(layer, "prefix", "")) + config = _shape_config(prefix, k, logical_n) + if config is None: + return False + padded_n = _HC_DOWN_PADDED_N if prefix.endswith(_HC_DOWN_SUFFIX) else logical_n + if padded_n % 32 or k % 16: + raise RuntimeError(f"online QPN8 received an unaligned weight {(padded_n, k)}") + + if padded_n != logical_n: + padding = weight.new_zeros((padded_n - logical_n, k)) + weight_for_quant = torch.cat((weight, padding), dim=0) + else: + weight_for_quant = weight + + weight_f32 = weight_for_quant.float() + channel_scales = weight_f32.abs().amax(dim=1, keepdim=True).div_(448.0) + channel_scales = torch.where( + channel_scales == 0, torch.ones_like(channel_scales), channel_scales + ) + qweight = (weight_f32 / channel_scales).to(torch.float8_e4m3fn) + codes, scales = sm70_ops.fp8_qpn8_prepare_sm70( + qweight.contiguous(), channel_scales.contiguous() + ) + workspace = _workspace(weight) + + replace_parameter(layer, "weight", codes) + replace_parameter(layer, "weight_scale_inv", scales) + setattr(layer, _CODES_ATTR, layer.weight) + setattr(layer, _SCALES_ATTR, layer.weight_scale_inv) + layer._sm70_qwen4_exp_online_qpn8_k = k + layer._sm70_qwen4_exp_online_qpn8_n = padded_n + layer._sm70_qwen4_exp_online_qpn8_logical_n = logical_n + layer._sm70_qwen4_exp_online_qpn8_split_k = config[0] + layer._sm70_qwen4_exp_online_qpn8_nacc = config[1] + layer._sm70_qwen4_exp_online_qpn8_prefetch = config[2] + layer._sm70_qwen4_exp_online_qpn8_workspace = workspace + layer._sm70_qwen4_exp_online_qpn8_workspace_ptr = workspace.data_ptr() + if prefix.endswith(_HC_DOWN_SUFFIX): + # Resolve the shared scratch tensor while weights are prepared. Looking + # it up through the process-global dictionary from forward makes + # TorchDynamo emit a dictionary clear/update, which cudagraph rejects + # as a module-state mutation. + layer.register_buffer( + _HC_PARTIALS_ATTR, + _hc_workspace(weight), + persistent=False, + ) + layer.sm70_fp8_qpn8 = True + layer.sm70_fp8_prefill_exact_dense_workspace_ptr = workspace.data_ptr() + setattr(layer, _STATE_ATTR, True) + logger.info_once( + "SM70 Qwen4Exp online channel-QPN8 enabled for selected FP16 decode " + "projections." + ) + return True + + +def maybe_apply_online_qpn8( + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor | None: + """Run a prepared channel-QPN8 projection, or return ``None``.""" + if not getattr(layer, _STATE_ATTR, False): + return None + k = int(layer._sm70_qwen4_exp_online_qpn8_k) + n = int(layer._sm70_qwen4_exp_online_qpn8_n) + logical_n = int(layer._sm70_qwen4_exp_online_qpn8_logical_n) + if x.dtype != torch.float16 or x.shape[-1] != k: + raise RuntimeError( + "SM70 Qwen4Exp online QPN8 requires FP16 inputs with " + f"K={k}, got dtype={x.dtype} shape={tuple(x.shape)}" + ) + x_2d = x.reshape(-1, k) + if not x_2d.is_contiguous(): + x_2d = x_2d.contiguous() + out = torch.empty((x_2d.size(0), n), dtype=x.dtype, device=x.device) + sm70_ops.fp8_qpn8_dispatch_sm70_out( + out, + int(layer._sm70_qwen4_exp_online_qpn8_workspace_ptr), + x_2d, + getattr(layer, _CODES_ATTR), + getattr(layer, _SCALES_ATTR), + int(layer._sm70_qwen4_exp_online_qpn8_split_k), + int(layer._sm70_qwen4_exp_online_qpn8_nacc), + bool(layer._sm70_qwen4_exp_online_qpn8_prefetch), + False, + ) + logical_out = out[:, :logical_n] + if bias is not None: + logical_out.add_(bias) + return logical_out.reshape(*x.shape[:-1], logical_n) + + +def maybe_apply_fused_hc( + down_layer: nn.Module, + up_layer: nn.Module, + xn: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Run the exact Qwen4Exp HC projection pair through one opaque route.""" + if not ( + getattr(down_layer, _STATE_ATTR, False) + and getattr(up_layer, _STATE_ATTR, False) + ): + return None + if not sm70_ops.has_fp8_qpn8_hc_dispatch(): + raise RuntimeError("prepared Qwen4Exp HC requires the fused QPN8 op") + if xn.dtype != torch.float16 or xn.ndim != 2 or xn.size(1) != 10240: + raise RuntimeError( + "fused Qwen4Exp HC requires a contiguous FP16 [M, 10240] input" + ) + if not xn.is_contiguous(): + xn = xn.contiguous() + if ( + int(down_layer._sm70_qwen4_exp_online_qpn8_k) != 10240 + or int(down_layer._sm70_qwen4_exp_online_qpn8_n) != _HC_DOWN_PADDED_N + or int(up_layer._sm70_qwen4_exp_online_qpn8_k) != 320 + or int(up_layer._sm70_qwen4_exp_online_qpn8_n) != 10240 + ): + raise RuntimeError("fused Qwen4Exp HC received incompatible weights") + + down_workspace_ptr = int(down_layer._sm70_qwen4_exp_online_qpn8_workspace_ptr) + up_workspace_ptr = int(up_layer._sm70_qwen4_exp_online_qpn8_workspace_ptr) + if down_workspace_ptr != up_workspace_ptr: + raise RuntimeError("fused Qwen4Exp HC requires one shared dense workspace") + partials = getattr(down_layer, _HC_PARTIALS_ATTR, None) + if not isinstance(partials, torch.Tensor): + raise RuntimeError("fused Qwen4Exp HC requires prepared partials") + + m = xn.size(0) + block_out = xn.new_empty((m, 2560)) + injection_out = xn.new_empty((m, 4)) + down_staging = xn.new_empty((m, _HC_DOWN_PADDED_N)) + lora_staging = xn.new_empty((m, 320)) + gate_staging = xn.new_empty((m, 10240)) + sm70_ops.fp8_qpn8_hc_dispatch_sm70_out( + block_out, + injection_out, + down_staging, + lora_staging, + gate_staging, + partials, + down_workspace_ptr, + xn, + getattr(down_layer, _CODES_ATTR), + getattr(down_layer, _SCALES_ATTR), + getattr(up_layer, _CODES_ATTR), + getattr(up_layer, _SCALES_ATTR), + ) + return block_out, injection_out + + +__all__ = [ + "_shape_config", + "maybe_apply_fused_hc", + "maybe_apply_online_qpn8", + "maybe_prepare_online_qpn8", +] diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index f9152a52e7..68da508afc 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -28,10 +28,7 @@ def is_fp4_marlin_supported(): return False if current_platform.has_device_capability(75): return True - return ( - current_platform.is_device_capability((7, 0)) - and ops.sm70_marlin_available() - ) + return current_platform.is_device_capability((7, 0)) and ops.sm70_marlin_available() def _nvfp4_compute_scale_factor( diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py index d10ef0f830..63ae680d70 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py @@ -21,6 +21,7 @@ val=torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32) ) _TRITON_NVFP4_EMULATION_SUPPORTED: bool | None = None +_TRITON_NVFP4_DEQUANTIZATION_SUPPORTED: bool | None = None def _supports_triton_nvfp4_emulation(device: torch.device | None = None) -> bool: @@ -36,6 +37,27 @@ def _supports_triton_nvfp4_emulation(device: torch.device | None = None) -> bool return _TRITON_NVFP4_EMULATION_SUPPORTED +def _supports_triton_nvfp4_dequantization( + device: torch.device | None = None, +) -> bool: + """Return whether the weight-only Triton dequantizer can run. + + The full quantize-dequantize emulation kernel needs the native + ``tl.float8e4nv`` type and therefore remains gated to SM89+. Weight + dequantization can also run on older CUDA GPUs because its E4M3 scale + bytes are decoded with regular integer and float32 operations. + """ + if device is not None and device.type != "cuda": + return False + + global _TRITON_NVFP4_DEQUANTIZATION_SUPPORTED + if _TRITON_NVFP4_DEQUANTIZATION_SUPPORTED is None: + _TRITON_NVFP4_DEQUANTIZATION_SUPPORTED = _supports_triton_nvfp4_emulation( + device + ) or (current_platform.is_cuda() and current_platform.has_device_capability(70)) + return _TRITON_NVFP4_DEQUANTIZATION_SUPPORTED + + def _quantize_e4m3fn_scale(scale: torch.Tensor) -> torch.Tensor: if not scale.is_cuda or _supports_triton_nvfp4_emulation(scale.device): return scale.to(torch.float8_e4m3fn).to(torch.float32) @@ -77,6 +99,31 @@ def _e2m1_inline(magnitude): return tl.where(b2 == 1, high_group, low_group) +@triton.jit +def _apply_float32_sign_bit(value, sign_bit): + """Apply a sign bit without canonicalizing negative zero.""" + value_bits = tl.cast(value, tl.uint32, bitcast=True) + signed_bits = value_bits | (sign_bit.to(tl.uint32) << 31) + return tl.cast(signed_bits, tl.float32, bitcast=True) + + +@triton.jit +def _e4m3fn_byte_to_float(raw): + """Decode E4M3FN bytes without using a native FP8 Triton type.""" + raw_i32 = raw.to(tl.int32) + sign_bit = (raw_i32 >> 7) & 1 + exponent = (raw_i32 >> 3) & 0x0F + mantissa = raw_i32 & 0x07 + mantissa_f32 = mantissa.to(tl.float32) + normal = (1.0 + mantissa_f32 * 0.125) * tl.exp2(exponent.to(tl.float32) - 7.0) + subnormal = mantissa_f32 * 0.001953125 # 2**-9 + value = _apply_float32_sign_bit( + tl.where(exponent == 0, subnormal, normal), sign_bit + ) + is_nan = (exponent == 0x0F) & (mantissa == 0x07) + return tl.where(is_nan, float("nan"), value) + + @triton.jit def _dequantize_nvfp4_kernel( fp4_ptr, @@ -88,6 +135,7 @@ def _dequantize_nvfp4_kernel( BLOCK_SIZE: tl.constexpr, has_batch_global_scale: tl.constexpr, TILE_BLOCKS: tl.constexpr, + NATIVE_FP8_SCALE_DECODE: tl.constexpr, ): """Triton kernel for NVFP4 dequantization (swizzle=False). @@ -119,7 +167,10 @@ def _dequantize_nvfp4_kernel( mask=block_mask, other=0, ) - scale_f32 = tl.cast(raw_scales, tl.float8e4nv, bitcast=True).to(tl.float32) + if NATIVE_FP8_SCALE_DECODE: + scale_f32 = tl.cast(raw_scales, tl.float8e4nv, bitcast=True).to(tl.float32) + else: + scale_f32 = _e4m3fn_byte_to_float(raw_scales) scale_values = (scale_f32 * global_scale)[:, None] # Load [TILE_BLOCKS, BLOCK_PACKED] packed bytes @@ -139,12 +190,20 @@ def _dequantize_nvfp4_kernel( low_mag = low_nibble & 0x07 low_val = _e2m1_inline(low_mag) low_sign = (low_nibble >> 3) & 1 - low_result = tl.where(low_sign == 1, -low_val, low_val) * scale_values + if NATIVE_FP8_SCALE_DECODE: + low_signed = tl.where(low_sign == 1, -low_val, low_val) + else: + low_signed = _apply_float32_sign_bit(low_val, low_sign) + low_result = low_signed * scale_values high_mag = high_nibble & 0x07 high_val = _e2m1_inline(high_mag) high_sign = (high_nibble >> 3) & 1 - high_result = tl.where(high_sign == 1, -high_val, high_val) * scale_values + if NATIVE_FP8_SCALE_DECODE: + high_signed = tl.where(high_sign == 1, -high_val, high_val) + else: + high_signed = _apply_float32_sign_bit(high_val, high_sign) + high_result = high_signed * scale_values # Interleave for coalesced contiguous store result = tl.interleave(low_result, high_result) @@ -345,6 +404,7 @@ def _triton_dequantize_nvfp4( block_size, is_3d, tile_blocks, + _supports_triton_nvfp4_emulation(tensor_fp4.device), num_warps=nw, num_stages=ns, ) @@ -406,7 +466,7 @@ def dequantize_to_dtype( # Two fp4 values are packed into one uint8. assert tensor_fp4.dtype == torch.uint8 - if not swizzle and _supports_triton_nvfp4_emulation(tensor_fp4.device): + if not swizzle and _supports_triton_nvfp4_dequantization(tensor_fp4.device): return _triton_dequantize_nvfp4( tensor_fp4, tensor_sf, global_scale, dtype, block_size ) diff --git a/vllm/model_executor/layers/quantization/utils/w8a8_utils.py b/vllm/model_executor/layers/quantization/utils/w8a8_utils.py index 40a350f4d8..7a08f380d6 100644 --- a/vllm/model_executor/layers/quantization/utils/w8a8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/w8a8_utils.py @@ -22,8 +22,7 @@ def cutlass_fp8_supported() -> bool: capability_tuple = current_platform.get_device_capability() capability = -1 if capability_tuple is None else capability_tuple.to_int() - return _cutlass_capability_probe("cutlass_scaled_mm_supports_fp8", - capability) + return _cutlass_capability_probe("cutlass_scaled_mm_supports_fp8", capability) def cutlass_block_fp8_supported() -> bool: @@ -33,8 +32,7 @@ def cutlass_block_fp8_supported() -> bool: capability_tuple = current_platform.get_device_capability() capability = -1 if capability_tuple is None else capability_tuple.to_int() - return _cutlass_capability_probe("cutlass_scaled_mm_supports_block_fp8", - capability) + return _cutlass_capability_probe("cutlass_scaled_mm_supports_block_fp8", capability) def cutlass_group_gemm_supported() -> bool: @@ -44,8 +42,7 @@ def cutlass_group_gemm_supported() -> bool: capability_tuple = current_platform.get_device_capability() capability = -1 if capability_tuple is None else capability_tuple.to_int() - return _cutlass_capability_probe("cutlass_group_gemm_supported", - capability) + return _cutlass_capability_probe("cutlass_group_gemm_supported", capability) CUTLASS_FP8_SUPPORTED = cutlass_fp8_supported() diff --git a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py b/vllm/model_executor/layers/sparse_attn_indexer_kpool.py new file mode 100644 index 0000000000..60ac0fe209 --- /dev/null +++ b/vllm/model_executor/layers/sparse_attn_indexer_kpool.py @@ -0,0 +1,1054 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Custom Sparse Attention Indexer layers.""" + +import os + +import torch + +import vllm.envs as envs +from vllm._aiter_ops import rocm_aiter_ops +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import get_current_vllm_config_or_none +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.custom_op import CustomOp +from vllm.models.glm5next.nvidia.ops.kpool_compress import ( + expand_pools_and_append_tail, + expand_pools_to_tokens, + kpool_compress_and_write_cache, + kpool_decode_update_and_maybe_write_cache_batched, + kpool_seed_tail_cache, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import ( + fp8_fp4_mqa_logits, + fp8_fp4_paged_mqa_logits, + has_deep_gemm, +) +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, +) +from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton +from vllm.v1.worker.workspace import current_workspace_manager + +if current_platform.is_cuda_alike(): + from vllm import _custom_ops as ops +elif current_platform.is_xpu(): + from vllm._xpu_ops import xpu_ops + +logger = init_logger(__name__) + +RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 + +# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32. +MXFP4_BLOCK_SIZE = 32 + + +def _is_exact_sm70_cuda() -> bool: + return current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) + + +# kpool write helper: form pools from the current token batch and compress them +# into the index K cache via the fused Triton kernel. + + +def _kpool_compress_insert( + k: torch.Tensor, + gate_score: torch.Tensor, + ape: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kpool: int, + head_dim: int, + round_scale: bool, +) -> None: + """Pool ``kpool`` consecutive tokens into one fp8 K and write at pool slots. + + ``slot_mapping`` is pool-granular (compress_ratio == kpool on the spec): + only the *last* token of each complete pool carries a valid (>=0) slot; + intra-pool tokens are -1. Every position is treated as a pool-completion + candidate and non-completions are masked off inside the kernel. Compacting + the valid rows first costs two device syncs on the eager prefill path and + buys nothing numerically. Assumes pool-aligned chunk starts. + """ + n = slot_mapping.shape[0] + # No pool can complete in a batch smaller than one pool; also keeps the + # clamped gather indices below in bounds. + if n < kpool: + return + pos = torch.arange(n, device=k.device) + valid = slot_mapping >= 0 + # Drop pools whose start falls before the batch (leading padding); their + # gate/k data is undefined anyway. + write_mask = valid & (pos >= kpool - 1) + offs = torch.arange(kpool, device=k.device) + idx = (pos - (kpool - 1)).clamp_min(0)[:, None] + offs[None, :] + kpool_compress_and_write_cache( + kv_cache, + k[idx], # [n, kpool, head_dim] + gate_score[idx], + ape, + slot_mapping.to(torch.int64), + pool_size=kpool, + head_dim=head_dim, + write_mask=write_mask, + round_scale=round_scale, + write_cache=True, + return_compressed=False, + ) + + +def _build_decode_scatter_indices( + decode_lens: torch.Tensor, + num_requests: int, + n: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-token (request id, intra-request index) for a non-uniform decode + batch, with ``n == decode_lens.sum()`` as a host int (avoids a + device sync and keeps both repeat_interleaves sync-free). + + Shared by every ``_scatter_decode_tokens_by_request`` call in a step: + building it per call would repeat the same repeat_interleave/cumsum chain + up to 5x per layer on the eager decode break. + """ + device = decode_lens.device + dl = decode_lens.to(torch.int64) + req_id = torch.repeat_interleave( + torch.arange(num_requests, device=device, dtype=torch.int64), + dl, + output_size=n, + ) + req_starts = torch.cumsum( + torch.cat([torch.zeros(1, device=device, dtype=torch.int64), dl[:-1]]), + dim=0, + ) + # Broadcast the per-request start offsets to per-token (length n == + # dl.sum()) so each token's intra-request index subtracts its own + # request's start. + starts = torch.repeat_interleave(req_starts, dl, output_size=n) + intra = torch.arange(n, device=device, dtype=torch.int64) - starts + return req_id, intra + + +def _scatter_decode_tokens_by_request( + tokens: torch.Tensor, + pad_value, + num_requests: int, + lmax: int, + scatter_indices: tuple[torch.Tensor, torch.Tensor], +) -> torch.Tensor: + """Group ``[N, ...]`` decode tokens into a padded ``[num_requests, lmax, ...]`` + layout: request ``r``'s tokens at row ``r`` in order; short requests padded. + + Unlike ``pack_seq_triton`` this is dtype-agnostic (needed for the int32 + slot/pos tensors) — it scatters with the shared per-step indices from + ``_build_decode_scatter_indices``. Used only for the non-uniform + (``requires_padding``) decode batch; uniform batches use a zero-copy + reshape. + """ + req_id, intra = scatter_indices + out = torch.full( + (num_requests, lmax, *tokens.shape[1:]), + pad_value, + dtype=tokens.dtype, + device=tokens.device, + ) + out[req_id, intra] = tokens + return out + + +def _decode_topk_seq_lens( + positions: torch.Tensor, + decode_lens: torch.Tensor, + num_decode_tokens: int, + batch_size: int, + next_n: int, + requires_padding: bool, +) -> torch.Tensor: + """Token-granular seq_len (pos + 1) per pool-topk row, layout-aware. + + ``pool_topk`` (and the logits it comes from) follow the padded + ``[batch_size, next_n]`` grid whenever ``requires_padding`` is set, so row + ``(b, t)`` corresponds to flat decode token ``offset_b + t`` -- NOT + ``b * next_n + t``. Slicing flat ``positions[: batch_size * next_n]`` + (the uniform-layout shortcut) misaligns every row after the first + non-uniform request and, past the decode region, reads prefill tokens' + positions; ``expand_pools_and_append_tail`` then anchors the tail at + another request's length, dropping the row's real tail tokens or emitting + indices past its sequence (out-of-bounds block-table reads). Padded rows + get 0 (empty tail); they are dropped by ``unpack_seq_triton`` anyway. + """ + n = batch_size * next_n + if not requires_padding: + return positions[:n].to(torch.int32) + 1 + scatter_idx = _build_decode_scatter_indices( + decode_lens, batch_size, num_decode_tokens + ) + padded = _scatter_decode_tokens_by_request( + positions[:num_decode_tokens].to(torch.int32), + -1, + batch_size, + next_n, + scatter_idx, + ) + return padded.reshape(n) + 1 # pad rows: -1 + 1 = 0 -> empty tail + + +def _gather_workspace_shapes( + total_seq_lens: int, + head_dim: int, + fp8_dtype: torch.dtype, + use_fp4_cache: bool, +) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]: + """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for + the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 + scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + + (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.""" + if use_fp4_cache: + return ( + ((total_seq_lens, head_dim // 2), torch.uint8), + ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8), + ) + return ( + ((total_seq_lens, head_dim), fp8_dtype), + ((total_seq_lens, 4), torch.uint8), + ) + + +def kv_cache_as_quant_view( + kv_cache: torch.Tensor, + head_dim: int, + use_fp4_cache: bool, +) -> torch.Tensor: + """4D ``[num_blocks, block_size, 1, head_width]`` view expected by + DeepGEMM, from the 3D indexer kv-cache allocation.""" + if use_fp4_cache: + assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 + num_blocks, block_size, _ = kv_cache.shape + page_bytes = int(kv_cache.stride(0)) + fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE + return torch.as_strided( + kv_cache, + size=(num_blocks, block_size, 1, fp4_bytes), + stride=(page_bytes, fp4_bytes, fp4_bytes, 1), + ) + return kv_cache.unsqueeze(-2) + + +@eager_break_during_capture +def sparse_attn_indexer_kpool( + hidden_states: torch.Tensor, + k_cache_prefix: LayerNameType, + kv_cache: torch.Tensor, + q_quant: torch.Tensor, + q_scale: torch.Tensor | None, + k: torch.Tensor, + weights: torch.Tensor, + quant_block_size: int, + scale_fmt: str | None, + topk_tokens: int, + head_dim: int, + max_model_len: int, + total_seq_lens: int, + topk_indices_buffer: torch.Tensor, + skip_k_cache_insert: bool, + use_fp4_cache: bool = False, + # kpool params (Plan-A: gate is consumed at write time and read back at + # topk time to softmax-weight the pool). + gate_score: torch.Tensor | None = None, + compress_ape: torch.Tensor | None = None, + index_kpool: int = 1, + positions: torch.Tensor | None = None, + # Paged tail cache (in-progress pool's raw K + gate score), replacing the + # transient _DECODE_TAIL ring. tail_prefix resolves attn_metadata[tail_prefix] + # for the tail group's token-granular slot_mapping. None on the dummy/profiling + # path and when the tail cache is disabled. + tail_kv_cache: torch.Tensor | None = None, + tail_prefix: str | None = None, +) -> torch.Tensor: + # careful! this will be None in dummy run + attn_metadata = get_forward_context().attn_metadata + fp8_dtype = current_platform.fp8_dtype() + k_cache_prefix = _resolve_layer_name(k_cache_prefix) + + # assert isinstance(attn_metadata, dict) + if not isinstance(attn_metadata, dict): + # Reserve workspace for indexer during profiling run + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache + ) + current_workspace_manager().get_simultaneous( + values_spec, + scales_spec, + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + ) + + # Reserve profiler-visible memory for the worst-case decode logits, + # whose shape is [B * next_n, max_model_len]. This profiling branch + # returns before invoking the logits kernel itself. + cfg = get_current_vllm_config_or_none() + worst_decode_tokens = 0 + if cfg is not None: + sched = cfg.scheduler_config + num_spec = ( + cfg.speculative_config.num_speculative_tokens + if cfg.speculative_config is not None + else 0 + ) + worst_decode_tokens = min( + sched.max_num_seqs * (num_spec + 1), + sched.max_num_batched_tokens, + ) + # float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes. + decode_logits_elems = worst_decode_tokens * max_model_len * 4 + prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + max_logits_elems = max(decode_logits_elems, prefill_cap_elems) + _ = torch.empty( + max_logits_elems, dtype=torch.uint8, device=hidden_states.device + ) + + return sparse_attn_indexer_kpool_fake( + hidden_states, + k_cache_prefix, + kv_cache, + q_quant, + q_scale, + k, + weights, + quant_block_size, + scale_fmt, + topk_tokens, + head_dim, + max_model_len, + total_seq_lens, + topk_indices_buffer, + skip_k_cache_insert, + use_fp4_cache, + ) + attn_metadata_narrowed = attn_metadata[k_cache_prefix] + assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata) + slot_mapping = attn_metadata_narrowed.slot_mapping + has_decode = attn_metadata_narrowed.num_decodes > 0 + has_prefill = attn_metadata_narrowed.num_prefills > 0 + num_decode_tokens = attn_metadata_narrowed.num_decode_tokens + sm70_fp16_indexer = _is_exact_sm70_cuda() + + # q_scale is required iff the FP4 cache path is enabled; the FP8 path + # folds the Q scale into `weights` inside fused_indexer_q_rope_quant. + if sm70_fp16_indexer: + assert q_quant.dtype == torch.float16 + assert q_scale is None + assert not use_fp4_cache, "SM70 requires the FP8 indexer cache" + elif use_fp4_cache: + assert q_scale is not None, "use_fp4_cache=True requires q_scale" + else: + assert q_scale is None, "q_scale must be None when use_fp4_cache=False" + + # During speculative decoding, k may be padded to the CUDA graph batch + # size while slot_mapping only covers actual tokens. Truncate k to avoid + # out-of-bounds reads in the kernel. + num_tokens = slot_mapping.shape[0] + if k is not None: + k = k[:num_tokens] + + if not skip_k_cache_insert: + assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet" + if index_kpool > 1 and gate_score is not None and compress_ape is not None: + # kpool prefill write: pool kpool consecutive prefill tokens via + # softmax(gate+ape)-weighted sum -> Hadamard -> fp8 -> pool slots. + # Decode tokens (the first num_decode_tokens in the batch) cannot be + # pooled here — their pool's earlier tokens are not in this batch — + # so they are deferred to the tail-buffer kernel in has_decode. + # compress_ratio == index_kpool makes slot_mapping pool-granular. + n_prefill = num_tokens - num_decode_tokens + if n_prefill > 0: + # decode tokens are batched first; prefill tokens follow. + prefill_slice = slice(num_decode_tokens, num_tokens) + _kpool_compress_insert( + k[prefill_slice], + gate_score[prefill_slice], + compress_ape, + kv_cache, + slot_mapping[prefill_slice], + index_kpool, + head_dim, + round_scale=(scale_fmt is not None), + ) + # Persist each request's incomplete prefill pool so decode can + # finish it, including after PD transfer. Tail slots use + # ``pos % kpool`` within the request's tail block. Processing + # only the batch's trailing tokens would miss all but the last + # request in a multi-request prefill. + if ( + tail_kv_cache is not None + and tail_prefix is not None + and os.environ.get("VLLM_KPOOL_SKIP_TAIL_CACHE") != "1" + ): + tail_meta = attn_metadata.get(_resolve_layer_name(tail_prefix)) + if tail_meta is not None: + assert isinstance(tail_meta, DeepseekV32IndexerMetadata) + kpool_seed_tail_cache( + tail_kv_cache, + k[prefill_slice], + gate_score[prefill_slice], + tail_meta.slot_mapping[prefill_slice], + index_kpool, + head_dim, + ) + else: + # standard: per-token fp8 quant + scatter (all tokens). + assert scale_fmt is not None + ops.indexer_k_quant_and_cache( + k, + kv_cache, + slot_mapping, + quant_block_size, + scale_fmt, + ) + + topk_indices_buffer[: hidden_states.shape[0]] = -1 + if has_prefill: + prefill_metadata = attn_metadata_narrowed.prefill + assert prefill_metadata is not None + + # Short sequences select every pool, so skip sparse scoring and fill + # the top-k buffer with all causal token indices. The index-K cache was + # already written above. + n_prefill_sf = num_tokens - num_decode_tokens + # Host-side short-prefill predicate: max_prefill_seq_len is computed + # in the metadata builder (exact for prefill rows) and equals + # positions[prefill_slice].max() + 1, so this replaces a + # positions.max().item() device sync per layer. -1 (unknown metadata) + # falls back to the device-side check. + if prefill_metadata.max_prefill_seq_len >= 0: + short_prefill = ( + n_prefill_sf > 0 + and positions is not None + and prefill_metadata.max_prefill_seq_len <= topk_tokens + ) + else: + short_prefill = ( + n_prefill_sf > 0 + and positions is not None + and int(positions[num_decode_tokens:num_tokens].max().item()) + 1 + <= topk_tokens + ) + if short_prefill: + # short_prefill is only True when positions is not None (above), + # but narrow explicitly for the indexer below. + assert positions is not None + _arange = torch.arange( + topk_indices_buffer.shape[1], + device=topk_indices_buffer.device, + dtype=torch.int32, + ) + _pos = positions[num_decode_tokens:num_tokens].to(torch.int32) + _buf = topk_indices_buffer[num_decode_tokens:num_tokens] + _buf[:] = _arange[None, :] + _buf[_arange[None, :] > _pos[:, None]] = -1 + + # Get the full shared workspace buffers once (will allocate on first use). + # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and + # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0 + # scales) based on use_fp4_cache. + workspace_manager = current_workspace_manager() + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache + ) + k_quant_full, k_scale_full = workspace_manager.get_simultaneous( + values_spec, + scales_spec, + ) + for chunk in prefill_metadata.chunks if not short_prefill else (): + k_quant = k_quant_full[: chunk.total_seq_lens] + k_scale = k_scale_full[: chunk.total_seq_lens] + + if not chunk.skip_kv_gather: + ops.cp_gather_indexer_k_quant_cache( + kv_cache, + k_quant, + k_scale, + chunk.block_table, + chunk.cu_seq_lens, + ) + + q_slice = q_quant[chunk.token_start : chunk.token_end] + q_scale_slice = ( + q_scale[chunk.token_start : chunk.token_end] + if q_scale is not None + else None + ) + # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 + # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. + if use_fp4_cache: + q_slice_cast = q_slice.view(torch.int8) + k_quant_cast = k_quant.view(torch.int8) + k_scale_cast = k_scale.view(torch.int32).squeeze(-1) + else: + q_slice_cast = q_slice + k_quant_cast = k_quant + k_scale_cast = k_scale.view(torch.float32).squeeze(-1) + if sm70_fp16_indexer: + from vllm.models.deepseek_v4.sm70.indexer import ( + sm70_indexer_prefill_logits, + ) + + logits = sm70_indexer_prefill_logits( + q_slice, + k_quant, + k_scale, + weights[chunk.token_start : chunk.token_end], + ) + else: + logits = fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + weights[chunk.token_start : chunk.token_end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + clean_logits=False, + ) + num_rows = logits.shape[0] + + # kpool: logits are pool-granular (compress_ratio == index_kpool), + # so topk selects pools. We pick topk_tokens // kpool pools then + # expand each pool back to its kpool constituent tokens. + select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens + if index_kpool > 1: + pool_topk = torch.empty( + (num_rows, select_k), dtype=torch.int32, device=logits.device + ) + topk_dst = pool_topk + else: + topk_dst = topk_indices_buffer[ + chunk.token_start : chunk.token_end, :topk_tokens + ] + + if current_platform.is_xpu(): + xpu_ops.top_k_per_row_prefill( # type: ignore[attr-defined] + logits, + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + topk_dst, + num_rows, + logits.stride(0), + logits.stride(1), + select_k, + ) + else: + torch.ops._C.top_k_per_row_prefill( + logits, + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + topk_dst, + num_rows, + logits.stride(0), + logits.stride(1), + select_k, + ) + + if index_kpool > 1: + pool_ids = pool_topk.to(torch.int64) + if positions is not None: + # Fused expand-pools + append-tail into one Triton kernel + # (replaces ~25 elementwise ops). seq_len is token-granular + # (pos+1); the kernel derives pool_len internally. + q_seq = ( + positions[chunk.token_start : chunk.token_end].to(torch.int32) + + 1 + ) + expanded = expand_pools_and_append_tail( + pool_ids, q_seq, index_kpool + ) + else: + valid = pool_ids >= 0 + expanded = expand_pools_to_tokens( + pool_ids, valid, topk_tokens, index_kpool + ) + topk_indices_buffer[ + chunk.token_start : chunk.token_end, : expanded.shape[-1] + ] = expanded + + if has_decode: + decode_metadata = attn_metadata_narrowed.decode + assert decode_metadata is not None + kv_cache_raw = kv_cache # raw [num_blocks, block_size, head_dim+4] for writes + if not sm70_fp16_indexer: + kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache) + + # Update the tail before reading logits; completed pools are compressed + # into the slot supplied by slot_mapping. + # Spec verification groups tokens by request and preserves position + # order so each token is stashed before the next completes its pool. + # Positions must remain token-granular because the kernel derives the + # pool phase and tail index from ``pos % kpool``. + if ( + index_kpool > 1 + and gate_score is not None + and compress_ape is not None + and positions is not None + and not skip_k_cache_insert + and os.environ.get("VLLM_KPOOL_SKIP_DECODE_WRITE") != "1" + ): + num_requests = attn_metadata_narrowed.num_decodes + # Kpool writes must recover the original request grouping after the + # indexer's flattened decode path. Host metadata avoids a CUDA graph + # sync when choosing the uniform or padded layout. + per_req_lens = decode_metadata.per_req_decode_lens + if per_req_lens is not None: + use_uniform = ( + decode_metadata.decode_is_uniform + and num_decode_tokens + == num_requests * decode_metadata.write_max_decode_len + ) + group_lens = per_req_lens + lmax = decode_metadata.write_max_decode_len + else: + # Legacy metadata without per-request lens: fall back to the + # host-side requires_padding flag. Unreached now (per-request + # lens is always populated for decode), kept defensive. + use_uniform = not decode_metadata.requires_padding + group_lens = decode_metadata.decode_lens + lmax = int(decode_metadata.decode_lens.max().item()) + if not use_uniform: + # Non-uniform decode_lens (mixed plain-decode + spec-verify, or + # a variable MTP-verify batch): scatter actual tokens into a + # padded [B, lmax] layout. int32 tensors can't go through + # pack_seq_triton (float/uint8 only). The scatter indices are + # shared by all five scatters below (and the tail slot one). + scatter_idx = _build_decode_scatter_indices( + group_lens, num_requests, num_decode_tokens + ) + dec_k = _scatter_decode_tokens_by_request( + k[:num_decode_tokens], 0, num_requests, lmax, scatter_idx + ) + dec_gate = _scatter_decode_tokens_by_request( + gate_score[:num_decode_tokens], + 0, + num_requests, + lmax, + scatter_idx, + ) + dec_slot = _scatter_decode_tokens_by_request( + slot_mapping[:num_decode_tokens], + -1, + num_requests, + lmax, + scatter_idx, + ) + dec_pos = _scatter_decode_tokens_by_request( + positions[:num_decode_tokens].to(torch.int32), + -1, + num_requests, + lmax, + scatter_idx, + ) + else: + next_n = num_decode_tokens // num_requests + shape2 = (num_requests, next_n) + dec_k = k[:num_decode_tokens].view(*shape2, head_dim) + dec_gate = gate_score[:num_decode_tokens].view(*shape2, head_dim) + dec_slot = slot_mapping[:num_decode_tokens].view(shape2) + dec_pos = positions[:num_decode_tokens].to(torch.int32).view(shape2) + tail_meta = ( + attn_metadata.get(_resolve_layer_name(tail_prefix)) + if tail_prefix is not None + else None + ) + # Paged tail cache replaces the transient _DECODE_TAIL ring. Group + # the tail group's token-granular slot_mapping per-request, mirroring + # dec_slot / dec_pos, so the kernel gets each request's current-token + # tail slot (block * kpool + pos % kpool). + if tail_meta is not None: + assert isinstance(tail_meta, DeepseekV32IndexerMetadata) + if tail_meta is None or tail_kv_cache is None: + dec_tail_slot = None + elif not use_uniform: + dec_tail_slot = _scatter_decode_tokens_by_request( + tail_meta.slot_mapping[:num_decode_tokens], + -1, + num_requests, + lmax, + scatter_idx, + ) + else: + dec_tail_slot = tail_meta.slot_mapping[:num_decode_tokens].view(shape2) + # The compress kernel writes the raw fp8 cache (not the quant view); + # pass the underlying kv_cache, not kv_cache_quant_view. + if dec_tail_slot is not None: + # Single batched launch over [num_requests, next_n] replaces the + # per-token sequential loop. The kernel iterates each request's + # tokens in position order internally, preserving the + # pool-completion read-after-stash dependency that the loop + # provided. Inputs are already grouped per request (uniform: + # view; non-uniform: _scatter_decode_tokens_by_request padded to + # [B, lmax]) — no per-token .contiguous() copies needed. + kpool_decode_update_and_maybe_write_cache_batched( + kv_cache_raw, + tail_kv_cache, + dec_tail_slot, + dec_k, + dec_gate, + compress_ape, + dec_slot, + dec_pos, + index_kpool, + head_dim, + round_scale=(scale_fmt is not None), + ) + decode_lens = decode_metadata.decode_lens + if decode_metadata.requires_padding: + # Padding also covers short chunked prefills classified as decode. + # MXFP4 uses zero-byte padding so padded slots dequantize to zero. + if q_scale is not None: + padded_q_quant_decode_tokens = pack_seq_triton( + q_quant[:num_decode_tokens], decode_lens, pad_value=0 + ) + padded_q_scale = pack_seq_triton( + q_scale[:num_decode_tokens], decode_lens, pad_value=0 + ) + else: + padded_q_quant_decode_tokens = pack_seq_triton( + q_quant[:num_decode_tokens], decode_lens + ) + padded_q_scale = None + padded_weights = pack_seq_triton( + weights[:num_decode_tokens], decode_lens, pad_value=0 + ).reshape(-1, *weights.shape[1:]) + else: + padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape( + decode_lens.shape[0], -1, *q_quant.shape[1:] + ) + if q_scale is not None: + padded_q_scale = q_scale[:num_decode_tokens].reshape( + decode_lens.shape[0], -1, *q_scale.shape[1:] + ) + else: + padded_q_scale = None + padded_weights = weights[:num_decode_tokens] + # TODO: move and optimize below logic with triton kernels + batch_size = padded_q_quant_decode_tokens.shape[0] + next_n = padded_q_quant_decode_tokens.shape[1] + num_padded_tokens = batch_size * next_n + seq_lens = decode_metadata.seq_lens[:batch_size] + # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1) + # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens; + # the downstream topk kernels accept both 1D and 2D. + padded_q_quant_cast = ( + padded_q_quant_decode_tokens.view(torch.int8) + if use_fp4_cache + else padded_q_quant_decode_tokens + ) + if sm70_fp16_indexer: + from vllm.models.deepseek_v4.sm70.indexer import ( + sm70_indexer_decode_logits, + ) + + logits = sm70_indexer_decode_logits( + padded_q_quant_decode_tokens.reshape( + num_padded_tokens, + padded_q_quant_decode_tokens.shape[-2], + padded_q_quant_decode_tokens.shape[-1], + ), + kv_cache_raw, + padded_weights[:num_padded_tokens].reshape(num_padded_tokens, -1), + seq_lens, + decode_metadata.block_table, + attn_metadata_narrowed.max_seq_len, + ) + else: + logits = fp8_fp4_paged_mqa_logits( + (padded_q_quant_cast, padded_q_scale), + kv_cache, + padded_weights[:num_padded_tokens], + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_model_len, + clean_logits=False, + ) + num_rows = logits.shape[0] + # kpool: logits are pool-granular -> select topk_tokens//kpool pools, + # then expand each pool back to its kpool tokens. + select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens + if index_kpool > 1: + pool_topk = torch.empty( + (num_rows, select_k), dtype=torch.int32, device=logits.device + ) + topk_dst = pool_topk + else: + topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens] + + if current_platform.is_cuda() and select_k in (512, 1024, 2048): + workspace_manager = current_workspace_manager() + (topk_workspace,) = workspace_manager.get_simultaneous( + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + ) + torch.ops._C.persistent_topk( + logits, + seq_lens, + topk_dst, + topk_workspace, + select_k, + attn_metadata_narrowed.max_seq_len, + ) + else: + if current_platform.is_xpu(): + xpu_ops.top_k_per_row_decode( # type: ignore[attr-defined] + logits, + next_n, + seq_lens, + topk_dst, + num_rows, + logits.stride(0), + logits.stride(1), + select_k, + ) + else: + torch.ops._C.top_k_per_row_decode( + logits, + next_n, + seq_lens, + topk_dst, + num_rows, + logits.stride(0), + logits.stride(1), + select_k, + ) + + # Resolve to token-level indices in the output buffer. + if index_kpool > 1: + pool_ids = pool_topk.to(torch.int64) + n = pool_topk.shape[0] + # Decode seq_lens are pool-granular; recover token lengths from + # positions using the padded [B, next_n] row layout when needed. + if positions is not None: + dec_seq = _decode_topk_seq_lens( + positions, + decode_lens, + num_decode_tokens, + batch_size, + next_n, + decode_metadata.requires_padding, + ) + else: + dec_seq = decode_metadata.seq_lens[:n] + if dec_seq.ndim == 2: + dec_seq = dec_seq[:, -1] + dec_seq = dec_seq.to(torch.int32) + out = expand_pools_and_append_tail(pool_ids, dec_seq, index_kpool) + else: + out = topk_dst + + if decode_metadata.requires_padding: + # Drop padded query rows introduced by the next_n padding above. + out = unpack_seq_triton( + out.reshape(batch_size, -1, out.shape[-1]), decode_lens + ) + topk_indices_buffer[: out.shape[0], : out.shape[-1]] = out + + return topk_indices_buffer + + +def sparse_attn_indexer_kpool_fake( + hidden_states: torch.Tensor, + k_cache_prefix: LayerNameType, + kv_cache: torch.Tensor, + q_quant: torch.Tensor, + q_scale: torch.Tensor | None, + k: torch.Tensor, + weights: torch.Tensor, + quant_block_size: int, + scale_fmt: str | None, + topk_tokens: int, + head_dim: int, + max_model_len: int, + total_seq_lens: int, + topk_indices_buffer: torch.Tensor | None, + skip_k_cache_insert: bool, + use_fp4_cache: bool = False, + gate_score: torch.Tensor | None = None, + compress_ape: torch.Tensor | None = None, + index_kpool: int = 1, + positions: torch.Tensor | None = None, + tail_kv_cache: torch.Tensor | None = None, + tail_prefix: str | None = None, +) -> torch.Tensor: + return topk_indices_buffer + + +direct_register_custom_op( + op_name="sparse_attn_indexer_kpool", + op_func=sparse_attn_indexer_kpool, + # The indexer writes the index-K cache in place (prefill k-cache insert + + # kpool decode write), so kv_cache must be declared as mutated — otherwise + # under full-graph compile dynamo assumes it is unchanged across the + # indexer→MLA boundary and the MLA reads stale/misaligned KV. The paged tail + # cache is likewise written in place (prefill tail scatter + decode stash). + mutates_args=["topk_indices_buffer", "kv_cache", "tail_kv_cache"], + fake_impl=sparse_attn_indexer_kpool_fake, + dispatch_key=current_platform.dispatch_key, +) + + +@CustomOp.register("sparse_attn_indexer_kpool") +class SparseAttnIndexerKpool(CustomOp): + """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a + separate custom op since it involves heavy custom kernels like `mqa_logits`, + `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires + specific memory layout or implementation for different hardware backends to + achieve optimal performance. + + For now, the default native path will use CUDA backend path. Other platform + may requires add the corresponding Custom Op name `sparse_attn_indexer` to + `custom_ops` in `CompilationConfig` to enable the platform specific path. + """ + + def __init__( + self, + k_cache, + quant_block_size: int, + scale_fmt: str, + topk_tokens: int, + head_dim: int, + max_model_len: int, + max_total_seq_len: int, + topk_indices_buffer: torch.Tensor, + skip_k_cache_insert: bool = False, + use_fp4_cache: bool = False, + tail_cache=None, + ): + super().__init__() + self.k_cache = k_cache + self.tail_cache = tail_cache + self.quant_block_size = quant_block_size + self.scale_fmt = scale_fmt + self.topk_tokens = topk_tokens + self.head_dim = head_dim + self.max_model_len = max_model_len + self.max_total_seq_len = max_total_seq_len + self.topk_indices_buffer = topk_indices_buffer + self.skip_k_cache_insert = skip_k_cache_insert + self.use_fp4_cache = use_fp4_cache + if ( + current_platform.is_cuda() + and not _is_exact_sm70_cuda() + and not has_deep_gemm() + ): + raise RuntimeError( + "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed." + ) + + def forward_native( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + weights: torch.Tensor, + *, + gate_score: torch.Tensor | None = None, + compress_ape: torch.Tensor | None = None, + index_kpool: int = 1, + positions: torch.Tensor | None = None, + ): + if current_platform.is_cuda() or current_platform.is_xpu(): + return self.forward_cuda( + hidden_states, + q_quant, + k, + weights, + gate_score=gate_score, + compress_ape=compress_ape, + index_kpool=index_kpool, + positions=positions, + ) + elif current_platform.is_rocm(): + return self.forward_hip(hidden_states, q_quant, k, weights) + else: + raise NotImplementedError( + "SparseAttnIndexer native forward is only implemented for " + "CUDA, ROCm and XPU platforms." + ) + + def forward_cuda( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + weights: torch.Tensor, + *, + gate_score: torch.Tensor | None = None, + compress_ape: torch.Tensor | None = None, + index_kpool: int = 1, + positions: torch.Tensor | None = None, + ): + # FP8 path: single tensor (per-token scale is folded into `weights`). + # FP4 path: (values, scales) tuple with scales required by the kernel. + if isinstance(q_quant, tuple): + q_values, q_scale = q_quant + else: + q_values, q_scale = q_quant, None + return torch.ops.vllm.sparse_attn_indexer_kpool( + hidden_states, + _encode_layer_name(self.k_cache.prefix), + self.k_cache.kv_cache, + q_values, + q_scale, + k, + weights, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + self.skip_k_cache_insert, + self.use_fp4_cache, + gate_score, + compress_ape, + index_kpool, + positions, + self.tail_cache.kv_cache if self.tail_cache is not None else None, + self.tail_cache.prefix if self.tail_cache is not None else None, + ) + + def forward_hip( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + weights: torch.Tensor, + ): + assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet" + assert isinstance(q_quant, torch.Tensor), ( + "AMD sparse_attn_indexer expects a single FP8 q_quant tensor" + ) + if rocm_aiter_ops.is_enabled(): + return torch.ops.vllm.rocm_aiter_sparse_attn_indexer( + hidden_states, + _encode_layer_name(self.k_cache.prefix), + self.k_cache.kv_cache, + q_quant, + k, + weights, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=self.skip_k_cache_insert, + ) + raise RuntimeError( + "Sparse attention indexer ROCm path is only supported on AITER. " + "Please enable aiter with VLLM_ROCM_USE_AITER=1" + ) diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 9f69084c62..e5f80ae2d5 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -65,6 +65,24 @@ def _sm70_dflash2_qpn8_rerank_requested() -> bool: ) +def _sm70_dflash2_use_dense_order() -> bool: + """Keep production on the scored full-vocabulary tie-order contract.""" + if envs.VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER: + return True + if envs.VLLM_SM70_DFLASH2_QPN8_ALLOW_CANDIDATE_ORDER: + logger.warning_once( + "Using experimental SM70 DFlash2 QPN8 candidate-order top-k. " + "This path is not authorized for quality-sensitive serving." + ) + return False + logger.warning_once( + "Ignoring VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER=0 without the explicit " + "benchmark-only VLLM_SM70_DFLASH2_QPN8_ALLOW_CANDIDATE_ORDER=1; " + "using the scored dense-order path." + ) + return True + + def _trace_sm70_lm_head_skip(reason: str) -> None: if envs.VLLM_SM70_GREEDY_TOKEN_FASTPATH_TRACE or envs.VLLM_SM70_PROFILE_TRACE: logger.warning_once("SM70 LM head fast path not prepared: %s", reason) @@ -537,7 +555,8 @@ def _maybe_sm70_dflash2_qpn8_rerank( values, _positions, ids = _sm70_dflash2_rerank_output_buffers( layer, num_rows, selector_k ) - if envs.VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER: + use_dense_order = _sm70_dflash2_use_dense_order() + if use_dense_order: _sm70_dflash2_dense_order_topk( layer._sm70_dflash2_rerank_dense_logits[:num_rows], qpn8_ids, @@ -595,7 +614,7 @@ def _maybe_sm70_dflash2_qpn8_rerank( logger.info_once( "SM70 DFlash2 QPN8 top-64 plus exact packed TurboMind FP16 rerank " "path enabled (dense_order=%s).", - envs.VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER, + use_dense_order, ) output_shape = (*x.shape[:-1], selector_k) return values.reshape(output_shape), ids.reshape(output_shape) @@ -824,6 +843,8 @@ def __init__( padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, quant_config: QuantizationConfig | None = None, prefix: str = "", + *, + quant_method: QuantizeMethodBase | None = None, ): super().__init__() self.prefix = prefix @@ -853,8 +874,10 @@ def __init__( ) self.embedding_dim = embedding_dim - quant_method = None - if quant_config is not None: + # Model-specific embeddings can preselect a storage method. This is + # required by Qwen4Exp PLE, whose table remains FP8 even when the + # routed experts use a different checkpoint quantization config. + if quant_method is None and quant_config is not None: quant_method = quant_config.get_quant_method(self, prefix=prefix) if quant_method is None: quant_method = UnquantizedEmbeddingMethod() @@ -1077,6 +1100,16 @@ def forward(self, input_): output_parallel = self.quant_method.embedding(self, masked_input.long()) # Mask the output embedding. if self.tp_size > 1: + if output_parallel.dtype in ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ): + # Each token row has exactly one TP owner. Communicate the raw + # FP8 bytes as int8 because NCCL does not reduce FP8 directly. + comm_output = output_parallel.view(torch.int8) + comm_output.masked_fill_(input_mask.unsqueeze(-1), 0) + output = tensor_model_parallel_all_reduce(comm_output) + return output.view(output_parallel.dtype) output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) # Reduce across all the model parallel GPUs. output = tensor_model_parallel_all_reduce(output_parallel) diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 2a5f746d78..3b4fcfa042 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -139,7 +139,8 @@ def device_loading_context(module: torch.nn.Module, target_device: torch.device) # Store original device states and move parameters to GPU if they're on CPU for name, p in module.named_parameters(): - if p.device.type == "cpu": + keep_on_cpu = getattr(p, "_vllm_keep_on_cpu", False) + if p.device.type == "cpu" and not keep_on_cpu: original_device_states[name] = p.device p.data = p.data.to(target_device) if getattr(p, "_vllm_is_uva_offloaded", False): diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index d48ec981af..1cb631a8af 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -361,9 +361,7 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: ) else: cache_config.mamba_cache_mode = ( - "all" - if model_config.supports_mamba_prefix_caching - else "align" + "all" if model_config.supports_mamba_prefix_caching else "align" ) logger.warning( "Mamba cache mode is set to '%s' for %s by default " @@ -570,6 +568,84 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: ) +def _strip_qwen4_exp_mrope(model_config: "ModelConfig") -> None: + configs = { + id(config): config + for config in ( + getattr(model_config, "hf_config", None), + model_config.hf_text_config, + ) + if config is not None + } + for config in configs.values(): + rope_parameters = getattr(config, "rope_parameters", None) + if rope_parameters is not None: + rope_parameters.pop("mrope_section", None) + rope_parameters.pop("mrope_interleaved", None) + + +class Qwen4ExpForConditionalGenerationConfig(Qwen3_5ForConditionalGenerationConfig): + """Apply the Qwen3.5 hybrid-cache contract to Qwen4Exp.""" + + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(vllm_config) + text_config = vllm_config.model_config.hf_text_config + if text_config.hc_count <= 1: + raise ValueError("Qwen4Exp requires hc_count > 1") + + parallel_config = vllm_config.parallel_config + uses_ple_or_qsa = bool(text_config.ple_layer_ids) or ( + getattr(text_config, "indexer_n_heads", None) is not None + ) + if uses_ple_or_qsa and ( + parallel_config.enable_dbo or parallel_config.ubatch_size > 1 + ): + raise NotImplementedError( + "Qwen4Exp PLE/QSA does not support dual-batch overlap or microbatching" + ) + + model_config = vllm_config.model_config + multimodal_config = model_config.multimodal_config + if multimodal_config is not None: + if not multimodal_config.language_model_only: + raise NotImplementedError( + "Qwen4Exp multimodal inference is not enabled in the initial " + "SM70 route; pass --language-model-only while the vision " + "tower remains outside the TP4 memory and quality gates" + ) + _strip_qwen4_exp_mrope(model_config) + + spec_config = vllm_config.speculative_config + if spec_config is not None and spec_config.method not in { + "mtp", + "ngram", + "ngram_gpu", + }: + raise NotImplementedError( + "Qwen4Exp speculative decoding supports only its native MTP " + "checkpoint and linear n-gram proposers" + ) + + +class Qwen4ExpForCausalLMConfig(Qwen4ExpForConditionalGenerationConfig): + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + _strip_qwen4_exp_mrope(vllm_config.model_config) + + +class Qwen4ExpMTPConfig(Qwen4ExpForConditionalGenerationConfig): + """Preserve MRoPE for a VL target and use 1D RoPE for a text target.""" + + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + if hasattr(vllm_config.model_config.hf_config, "vision_config"): + return + _strip_qwen4_exp_mrope(vllm_config.model_config) + + class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -630,6 +706,9 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig, "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, + "Qwen4ExpForCausalLM": Qwen4ExpForCausalLMConfig, + "Qwen4ExpForConditionalGeneration": Qwen4ExpForConditionalGenerationConfig, + "Qwen4ExpMTP": Qwen4ExpMTPConfig, "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, "XLMRobertaModel": JinaRobertaModelConfig, } diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index b546040b74..5fc0ba512a 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -102,6 +102,19 @@ def _get_max_soft_tokens( return None, False +def _build_broadcast_additive_padding_mask( + padding_positions: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + """Build an SDPA key-padding mask without materializing query rows.""" + batch_size, num_patches = padding_positions.shape + return torch.zeros( + (batch_size, 1, 1, num_patches), + dtype=dtype, + device=padding_positions.device, + ).masked_fill_(padding_positions[:, None, None, :], float("-inf")) + + # --------------------------------------------------------------------------- # Input schema # --------------------------------------------------------------------------- @@ -948,6 +961,7 @@ class Gemma4ForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() + model_config = vllm_config.model_config config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config @@ -955,9 +969,27 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.quant_config = quant_config self.multimodal_config = multimodal_config + # SM70 memory-efficient SDPA accepts fp16/fp32, but not bf16. Gemma4 + # checkpoints can independently mark their HF towers as bf16 even when + # the vLLM model dtype is fp16, which otherwise forces quadratic math + # attention. Select the model dtype during construction so subsequent + # UVA tower registration and weight loading preserve their storage. + self._use_sm70_memory_efficient_mm_sdpa = ( + current_platform.is_cuda() + and current_platform.is_device_capability(70) + and model_config.dtype == torch.float16 + ) + tower_init_kwargs = ( + {"dtype": model_config.dtype} + if self._use_sm70_memory_efficient_mm_sdpa + else {} + ) + # ---- Vision tower (shared by image and video) ---- with self._mark_tower_model(vllm_config, {"image", "video"}): - self.vision_tower = AutoModel.from_config(config=config.vision_config) + self.vision_tower = AutoModel.from_config( + config=config.vision_config, **tower_init_kwargs + ) self.embed_vision = Gemma4MultimodalEmbedder( config.vision_config, config.text_config ) @@ -965,7 +997,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # ---- Audio tower (variants with audio_config) ---- if config.audio_config is not None: with self._mark_tower_model(vllm_config, "audio"): - self.audio_tower = AutoModel.from_config(config=config.audio_config) + self.audio_tower = AutoModel.from_config( + config=config.audio_config, **tower_init_kwargs + ) # AutoModel.from_config does NOT call post_init(), # which is needed to initialize buffers that are absent # from the checkpoint (e.g. inv_timescales for relative @@ -1194,9 +1228,16 @@ def _process_image_input( pad_tensor = (pp_tensor == -1).all(dim=-1) inputs_embeds = vt.patch_embedder(pv_tensor, pp_tensor, pad_tensor) + attention_mask = ( + _build_broadcast_additive_padding_mask( + pad_tensor, inputs_embeds.dtype + ) + if self._use_sm70_memory_efficient_mm_sdpa + else ~pad_tensor + ) encoder_outputs = vt.encoder( inputs_embeds=inputs_embeds, - attention_mask=~pad_tensor, + attention_mask=attention_mask, pixel_position_ids=pp_tensor, ) hidden_states = encoder_outputs.last_hidden_state @@ -1302,9 +1343,14 @@ def _process_video_input( pad_chunk = padding_positions[i : i + max_batch_size] inputs_embeds = vt.patch_embedder(pv_chunk, pp_chunk, pad_chunk) + attention_mask = ( + _build_broadcast_additive_padding_mask(pad_chunk, inputs_embeds.dtype) + if self._use_sm70_memory_efficient_mm_sdpa + else ~pad_chunk + ) encoder_outputs = vt.encoder( inputs_embeds=inputs_embeds, - attention_mask=~pad_chunk, + attention_mask=attention_mask, pixel_position_ids=pp_chunk, ) last_hidden_states_list.append(encoder_outputs.last_hidden_state) diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index d6011a3e58..741b34fd05 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -32,11 +32,15 @@ from vllm.config import ModelConfig, SpeechToTextConfig, SpeechToTextParams from vllm.inputs import PromptType, TokensPrompt from vllm.logger import init_logger -from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncsByType, +) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.tasks import ScoreType from vllm.utils.collection_utils import common_prefix from vllm.utils.func_utils import supports_kw +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from .interfaces_base import VllmModel @@ -837,6 +841,14 @@ def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]: """ ... + @classmethod + def get_mamba_state_copy_funcs( + cls, + mamba_types: set[MambaAttentionBackendEnum], + ) -> MambaStateCopyFuncsByType: + copy_funcs = cls.get_mamba_state_copy_func() + return {mamba_type: copy_funcs for mamba_type in mamba_types} + @overload def is_hybrid(model: object) -> TypeIs[IsHybrid]: ... diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index f27c4f3037..31ab7968a8 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -34,6 +34,7 @@ from torch import nn from transformers import Qwen2MoeConfig +import vllm.envs as envs from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size @@ -142,6 +143,16 @@ def __init__( enforce_enable=_sm70_force_shared_expert_silu_custom_op(prefix) ) self.expert_gate = expert_gate + self._sm70_fused_shared_expert_gate = ( + envs.VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION + and expert_gate is not None + and prefix.endswith(".mlp.shared_expert") + and hidden_size == 2560 + and intermediate_size == 160 + and not reduce_results + and _sm70_force_shared_expert_silu_custom_op(prefix) + and hasattr(torch.ops._C, "sm70_f16_gate_mul_out") + ) def forward(self, x): x = _sm70_dump_qwen_mlp_tensor("mlp_input", self.layer_idx, x) @@ -149,15 +160,32 @@ def forward(self, x): out = fused_act(x) if fused_act is not None else None if out is None: gate_up, _ = self.gate_up_proj(x) - gate_up = _sm70_dump_qwen_mlp_tensor( - "mlp_gate_up", self.layer_idx, gate_up - ) + gate_up = _sm70_dump_qwen_mlp_tensor("mlp_gate_up", self.layer_idx, gate_up) out = self.act_fn(gate_up) out = _sm70_dump_qwen_mlp_tensor("mlp_silu_out", self.layer_idx, out) out, _ = self.down_proj(out) out = _sm70_dump_qwen_mlp_tensor("mlp_down_out", self.layer_idx, out) - if self.expert_gate is not None: + if ( + self._sm70_fused_shared_expert_gate + and x.shape[0] == 1 + and x.dtype == torch.float16 + and out.dtype == torch.float16 + ): + from vllm import _sm70_ops as sm70_ops + + assert self.expert_gate is not None + gate_weight = self.expert_gate.weight + if gate_weight.dtype != torch.float16 or not gate_weight.is_contiguous(): + raise RuntimeError( + "SM70 Qwen3Next fused shared-expert gate requires a " + "contiguous FP16 gate weight." + ) + sm70_ops.sm70_f16_gate_mul_out(out, x, gate_weight) + out = _sm70_dump_qwen_mlp_tensor( + "mlp_after_expert_gate", self.layer_idx, out + ) + elif self.expert_gate is not None: expert_gate = self.expert_gate(x)[0] expert_gate = _sm70_dump_qwen_mlp_tensor( "mlp_expert_gate", self.layer_idx, expert_gate diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index dea3dab7d8..88f6ba5ca1 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -462,8 +462,7 @@ def _mtp_quant_disabled_in_hf_config(*hf_configs: object) -> bool: ) if modules_to_not_convert and any( - str(module) in ("mtp", "model.mtp") - for module in modules_to_not_convert + str(module) in ("mtp", "model.mtp") for module in modules_to_not_convert ): return True diff --git a/vllm/model_executor/models/qwen3_dflash2.py b/vllm/model_executor/models/qwen3_dflash2.py index af6560385c..c2586108c5 100644 --- a/vllm/model_executor/models/qwen3_dflash2.py +++ b/vllm/model_executor/models/qwen3_dflash2.py @@ -12,6 +12,7 @@ from vllm.compilation.backends import set_model_tag from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig +from vllm.config.speculative import get_dflash_model_draft_tokens from vllm.distributed import ( get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, @@ -218,8 +219,9 @@ def __init__( hidden_size=config.hidden_size, taps=int(draft_config["conv_kernel_size"]), group_size=int(draft_config["conv_group_size"]), - # Query tokens per request: the bonus token plus the mask tokens. - block_size=1 + speculative_config.num_speculative_tokens, + # Keep the convolution boundary at the checkpoint-trained block + # even when lookup augments the target verification block. + block_size=1 + get_dflash_model_draft_tokens(speculative_config), params_dtype=vllm_config.model_config.dtype, ) self.attention_conv = DFlashGroupedConv( diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 5723dbe285..fce4e4b0e2 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -117,11 +117,16 @@ "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"), "Gemma4ForCausalLM": ("gemma4", "Gemma4ForCausalLM"), "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"), + "Qwen4ExpForCausalLM": ( + "vllm.models.qwen4_exp", + "Qwen4ExpForCausalLM", + ), "GlmForCausalLM": ("glm", "GlmForCausalLM"), "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"), "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"), "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"), "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"), + "Glm5NextForCausalLM": ("vllm.models.glm5next", "Glm5NextForCausalLM"), "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"), "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"), "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"), @@ -410,6 +415,10 @@ "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"), "Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vMoeForConditionalGeneration"), "GlmOcrForConditionalGeneration": ("glm_ocr", "GlmOcrForConditionalGeneration"), + "Glm5NextForConditionalGeneration": ( + "vllm.models.glm5next", + "Glm5NextForConditionalGeneration", + ), "GraniteSpeechForConditionalGeneration": ( "granite_speech", "GraniteSpeechForConditionalGeneration", @@ -562,6 +571,10 @@ "qwen3_5", "Qwen3_5MoeForConditionalGeneration", ), + "Qwen4ExpForConditionalGeneration": ( + "vllm.models.qwen4_exp", + "Qwen4ExpForConditionalGeneration", + ), "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"), "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"), "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"), @@ -629,9 +642,11 @@ "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"), "Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"), "GlmOcrMTPModel": ("glm_ocr_mtp", "GlmOcrMTP"), + "Glm5NextMTPModel": ("vllm.models.glm5next", "Glm5NextMTP"), "MedusaModel": ("medusa", "Medusa"), "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"), "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"), + "Qwen4ExpMTP": ("vllm.models.qwen4_exp", "Qwen4ExpMTP"), "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"), "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"), "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"), diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 095d0e363d..561bd419a4 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -46,18 +46,32 @@ class WeightsMapper: orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict) orig_to_new_substr: Mapping[str, str | None] = field(default_factory=dict) + orig_to_new_stacked: Mapping[str, tuple[str, str | int | tuple[int, ...]]] = field( + default_factory=dict + ) orig_to_new_prefix: Mapping[str, str | None] = field(default_factory=dict) orig_to_new_suffix: Mapping[str, str | None] = field(default_factory=dict) def __or__(self, other: "WeightsMapper") -> "WeightsMapper": """Combine two `WeightsMapper`s by merging their mappings.""" return WeightsMapper( + orig_to_new_regex={**self.orig_to_new_regex, **other.orig_to_new_regex}, orig_to_new_substr={**self.orig_to_new_substr, **other.orig_to_new_substr}, + orig_to_new_stacked={ + **self.orig_to_new_stacked, + **other.orig_to_new_stacked, + }, orig_to_new_prefix={**self.orig_to_new_prefix, **other.orig_to_new_prefix}, orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix}, ) def _map_name(self, key: str) -> str | None: + result = self._map_name_with_shard(key) + return result[0] if result is not None else None + + def _map_name_with_shard( + self, key: str + ) -> tuple[str, str | int | tuple[int, ...] | None] | None: for pattern, new_key in self.orig_to_new_regex.items(): if pattern.search(key): if new_key is None: @@ -72,6 +86,12 @@ def _map_name(self, key: str) -> str | None: key = key.replace(substr, new_key, 1) + shard_id: str | int | tuple[int, ...] | None = None + for substr, (new_key, new_shard_id) in self.orig_to_new_stacked.items(): + if substr in key: + key = key.replace(substr, new_key, 1) + shard_id = new_shard_id + for prefix, new_key in self.orig_to_new_prefix.items(): if key.startswith(prefix): if new_key is None: @@ -86,16 +106,19 @@ def _map_name(self, key: str) -> str | None: key = new_key.join(key.rsplit(suffix, 1)) - return key + return key, shard_id def apply( self, weights: Iterable[tuple[str, torch.Tensor]] ) -> Iterable[tuple[str, torch.Tensor]]: - return ( - (out_name, data) - for name, data in weights - if (out_name := self._map_name(name)) is not None - ) + for name, data in weights: + result = self._map_name_with_shard(name) + if result is None: + continue + out_name, shard_id = result + if shard_id is not None: + data.shard_id = shard_id + yield out_name, data def apply_list(self, values: list[str]) -> list[str]: return [ @@ -220,7 +243,14 @@ def _load_param( ) weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, weight_data) + try: + weight_loader(param, weight_data) + except Exception as exc: + raise RuntimeError( + f"Error loading weight {weight_qualname!r} " + f"with checkpoint shape {tuple(weight_data.shape)} into " + f"parameter shape {tuple(param.shape)}" + ) from exc logger.debug("Loaded weight %s with shape %s", weight_qualname, param.shape) @@ -354,6 +384,62 @@ def load_weights( return autoloaded_weights +def maybe_fuse_shared_experts( + weights: Iterable[tuple[str, torch.Tensor]], + *, + n_routed_experts: int, + n_shared_experts: int, + ckpt_prefix: str = "mlp.shared_experts", + enabled: bool | None = None, +) -> Iterable[tuple[str, torch.Tensor]]: + """Route AITER fused-shared-expert weights into routed-expert slots. + + CUDA leaves the input stream unchanged. On ROCm, AITER can fuse shared + experts into the routed expert tensor; in that case split the checkpoint + tensor into the extra expert slots expected by ``RoutedExperts``. + """ + if enabled is None: + from vllm._aiter_ops import rocm_aiter_ops + + enabled = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + if not enabled: + yield from weights + return + + prefix = f"{ckpt_prefix}." + for name, loaded_weight in weights: + if prefix not in name: + yield name, loaded_weight + continue + + split_dim = 1 if ("down_proj.weight" in name and loaded_weight.ndim > 1) else 0 + total = loaded_weight.shape[split_dim] + if total % n_shared_experts != 0: + raise ValueError( + f"FSE shared-expert weight {name!r} has size {total} along axis " + f"{split_dim}, not divisible by " + f"n_shared_experts={n_shared_experts}." + ) + chunk = total // n_shared_experts + for shared_expert_idx in range(n_shared_experts): + chunk_slice = slice( + shared_expert_idx * chunk, (shared_expert_idx + 1) * chunk + ) + if loaded_weight.ndim == 1: + chunk_weight = loaded_weight[chunk_slice] + elif split_dim == 0: + chunk_weight = loaded_weight[chunk_slice, :] + else: + chunk_weight = loaded_weight[:, chunk_slice] + yield ( + name.replace( + prefix, + f"mlp.experts.{n_routed_experts + shared_expert_idx}.", + ), + chunk_weight, + ) + + def init_vllm_registered_model( vllm_config: VllmConfig, *, diff --git a/vllm/models/common/__init__.py b/vllm/models/common/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/common/ops/__init__.py b/vllm/models/common/ops/__init__.py new file mode 100644 index 0000000000..1e4cbe3079 --- /dev/null +++ b/vllm/models/common/ops/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ops shared across model implementations.""" + +from .fused_qk_rmsnorm import fused_q_kv_rmsnorm + +__all__ = ["fused_q_kv_rmsnorm"] diff --git a/vllm/models/common/ops/fused_qk_rmsnorm.py b/vllm/models/common/ops/fused_qk_rmsnorm.py new file mode 100644 index 0000000000..66d43bc9c6 --- /dev/null +++ b/vllm/models/common/ops/fused_qk_rmsnorm.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +@triton.jit +def _fused_q_kv_rmsnorm_kernel( + q_ptr, + q_out_ptr, + q_weight_ptr, + q_in_stride, + q_out_stride, + kv_ptr, + kv_out_ptr, + kv_weight_ptr, + kv_in_stride, + kv_out_stride, + eps, + Q_SIZE: tl.constexpr, + KV_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + pid_task = tl.program_id(1) + + if pid_task == 0: + SIZE = Q_SIZE + row_in = q_ptr + token_idx * q_in_stride + weight_ptr = q_weight_ptr + row_out = q_out_ptr + token_idx * q_out_stride + else: + SIZE = KV_SIZE + row_in = kv_ptr + token_idx * kv_in_stride + weight_ptr = kv_weight_ptr + row_out = kv_out_ptr + token_idx * kv_out_stride + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + block = tl.arange(0, BLOCK_SIZE) + mask = block < SIZE + x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32) + variance = tl.sum(x * x, axis=0) / SIZE + rrms = tl.rsqrt(variance + eps) + weight = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32) + tl.store( + row_out + block, + (x * rrms * weight).to(row_out.dtype.element_ty), + mask=mask, + ) + + +def fused_q_kv_rmsnorm( + q: torch.Tensor, + kv: torch.Tensor, + q_weight: torch.Tensor, + kv_weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + assert q.ndim == 2 and kv.ndim == 2 + assert q.shape[0] == kv.shape[0], f"token dim mismatch: q={q.shape}, kv={kv.shape}" + assert q.stride(-1) == 1 and kv.stride(-1) == 1 + assert q_weight.is_contiguous() and kv_weight.is_contiguous() + + q_out = torch.empty_like(q) + kv_out = torch.empty_like(kv) + if q.shape[0] == 0: + return q_out, kv_out + + block_size = triton.next_power_of_2(max(q.shape[1], kv.shape[1])) + _fused_q_kv_rmsnorm_kernel[(q.shape[0], 2)]( + q, + q_out, + q_weight, + q.stride(0), + q_out.stride(0), + kv, + kv_out, + kv_weight, + kv.stride(0), + kv_out.stride(0), + eps, + Q_SIZE=q.shape[1], + KV_SIZE=kv.shape[1], + BLOCK_SIZE=block_size, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return q_out, kv_out diff --git a/vllm/models/common/ops/sequence_parallel.py b/vllm/models/common/ops/sequence_parallel.py new file mode 100644 index 0000000000..8631405eae --- /dev/null +++ b/vllm/models/common/ops/sequence_parallel.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) + + +def _custom_collective(name: str, x: torch.Tensor) -> torch.Tensor | None: + device_communicator = get_tp_group().device_communicator + if device_communicator is None: + return None + collective = getattr(device_communicator, name, None) + return None if collective is None else collective(x) + + +def sp_all_gather(x: torch.Tensor) -> torch.Tensor: + output = _custom_collective("custom_all_gather", x) + if output is not None: + return output + return tensor_model_parallel_all_gather(x, 0) + + +def sp_reduce_scatter(x: torch.Tensor) -> torch.Tensor: + assert x.ndim == 2 + tp_size = get_tensor_model_parallel_world_size() + sp_pad = (-x.shape[0]) % tp_size + if sp_pad > 0: + x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad)) + output = _custom_collective("custom_reduce_scatter", x) + if output is not None: + return output + return tensor_model_parallel_reduce_scatter(x, 0) + + +def sp_shard(x: torch.Tensor) -> torch.Tensor: + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + sp_pad = (-x.shape[0]) % tp_size + if sp_pad > 0: + pad = (0, 0) * (x.ndim - 1) + (0, sp_pad) + x = torch.nn.functional.pad(x, pad) + chunk = x.shape[0] // tp_size + return x[tp_rank * chunk : (tp_rank + 1) * chunk] + + +def sp_padding_mask( + is_padding: torch.Tensor | None, + hidden_states: torch.Tensor, +) -> torch.Tensor: + num_tokens = hidden_states.shape[0] + if is_padding is None: + is_padding = hidden_states.new_zeros(num_tokens, dtype=torch.bool) + assert is_padding.shape[0] == num_tokens + + tp_size = get_tensor_model_parallel_world_size() + sp_pad = (-num_tokens) % tp_size + if sp_pad > 0: + is_padding = torch.nn.functional.pad(is_padding, (0, sp_pad), value=True) + chunk = is_padding.shape[0] // tp_size + tp_rank = get_tensor_model_parallel_rank() + return is_padding[tp_rank * chunk : (tp_rank + 1) * chunk] diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 728993463e..bbdd5b1053 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -416,6 +416,11 @@ def compressor_kv_score() -> torch.Tensor: hidden_states, compressor.fused_wkv_wgate.weight, torch.float32, + getattr( + compressor.fused_wkv_wgate, + "_sm70_dsv4_fp13_weight", + None, + ), ) if output is not None: return output @@ -435,6 +440,11 @@ def indexer_weights_proj() -> torch.Tensor: hidden_states, indexer.weights_proj.weight, hidden_states.dtype, + getattr( + indexer.weights_proj, + "_sm70_dsv4_fp13_weight", + None, + ), ) if output is not None: return output @@ -447,6 +457,11 @@ def indexer_compressor_kv_score() -> torch.Tensor: hidden_states, indexer.compressor.fused_wkv_wgate.weight, torch.float32, + getattr( + indexer.compressor.fused_wkv_wgate, + "_sm70_dsv4_fp13_weight", + None, + ), ) if output is not None: return output @@ -851,6 +866,7 @@ def __init__( quant_config=None, prefix=f"{prefix}.weights_proj", ) + self.weights_proj._sm70_dsv4_fp13_gemv = True self.softmax_scale = self.head_dim**-0.5 self.scale_fmt = "ue8m0" diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 70fd8e3ec3..f9d8d43e4c 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -304,6 +304,7 @@ def __init__( disable_tp=True, prefix=f"{prefix}.fused_wkv_wgate", ) + self.fused_wkv_wgate._sm70_dsv4_fp13_gemv = True self.norm = RMSNorm(self.head_dim, self.rms_norm_eps) self.state_cache = CompressorStateCache( diff --git a/vllm/models/deepseek_v4/sm70/gemv.py b/vllm/models/deepseek_v4/sm70/gemv.py index d1374fbb76..20378173e7 100644 --- a/vllm/models/deepseek_v4/sm70/gemv.py +++ b/vllm/models/deepseek_v4/sm70/gemv.py @@ -16,6 +16,11 @@ _K = 4096 _BLOCK_K = 1024 _NUM_WARPS = 4 +_FP13_GROUP_VALUES = 32 +_FP13_GROUP_WORDS = 13 +_FP13_GROUPS_PER_ROW = _K // _FP13_GROUP_VALUES +_FP13_WORDS_PER_ROW = _FP13_GROUPS_PER_ROW * _FP13_GROUP_WORDS +_FP13_BUFFER = "_sm70_dsv4_fp13_weight" @triton.jit @@ -36,15 +41,240 @@ def _sm70_dsv4_fp16_gemv_kernel( tl.store(out_ptr + row, acc) -def can_use_sm70_dsv4_fp16_gemv( +@triton.jit +def _load_fp16x8(ptr): + word0, word1, word2, word3 = tl.inline_asm_elementwise( + "ld.global.v4.u32 {$0, $1, $2, $3}, [$4];", + constraints="=r,=r,=r,=r,l,~{memory}", + args=[ptr], + dtype=(tl.uint32, tl.uint32, tl.uint32, tl.uint32), + is_pure=False, + pack=1, + ) + return ( + (word0 & 0xFFFF).to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32), + (word0 >> 16).to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32), + (word1 & 0xFFFF).to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32), + (word1 >> 16).to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32), + (word2 & 0xFFFF).to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32), + (word2 >> 16).to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32), + (word3 & 0xFFFF).to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32), + (word3 >> 16).to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32), + ) + + +@triton.jit +def _select_fp13_word( + word0, + word1, + word2, + word3, + word4, + word5, + word6, + word7, + word8, + word9, + word10, + word11, + word12, + INDEX: tl.constexpr, +): + if INDEX == 0: # noqa: SIM116 - Triton constexpr dispatch needs static branches. + return word0 + elif INDEX == 1: + return word1 + elif INDEX == 2: + return word2 + elif INDEX == 3: + return word3 + elif INDEX == 4: + return word4 + elif INDEX == 5: + return word5 + elif INDEX == 6: + return word6 + elif INDEX == 7: + return word7 + elif INDEX == 8: + return word8 + elif INDEX == 9: + return word9 + elif INDEX == 10: + return word10 + elif INDEX == 11: + return word11 + else: + return word12 + + +@triton.jit +def _sm70_dsv4_fp13_gemv_kernel( + x_ptr, + packed_ptr, + out_ptr, + GROUP_VALUES: tl.constexpr, + GROUP_WORDS: tl.constexpr, + GROUPS_PER_ROW: tl.constexpr, + WORDS_PER_ROW: tl.constexpr, +): + row = tl.program_id(0) + groups = tl.arange(0, GROUPS_PER_ROW) + word_base = row * WORDS_PER_ROW + groups * GROUP_WORDS + word0 = tl.load(packed_ptr + word_base).to(tl.uint32, bitcast=True) + word1 = tl.load(packed_ptr + word_base + 1).to(tl.uint32, bitcast=True) + word2 = tl.load(packed_ptr + word_base + 2).to(tl.uint32, bitcast=True) + word3 = tl.load(packed_ptr + word_base + 3).to(tl.uint32, bitcast=True) + word4 = tl.load(packed_ptr + word_base + 4).to(tl.uint32, bitcast=True) + word5 = tl.load(packed_ptr + word_base + 5).to(tl.uint32, bitcast=True) + word6 = tl.load(packed_ptr + word_base + 6).to(tl.uint32, bitcast=True) + word7 = tl.load(packed_ptr + word_base + 7).to(tl.uint32, bitcast=True) + word8 = tl.load(packed_ptr + word_base + 8).to(tl.uint32, bitcast=True) + word9 = tl.load(packed_ptr + word_base + 9).to(tl.uint32, bitcast=True) + word10 = tl.load(packed_ptr + word_base + 10).to(tl.uint32, bitcast=True) + word11 = tl.load(packed_ptr + word_base + 11).to(tl.uint32, bitcast=True) + word12 = tl.load(packed_ptr + word_base + 12).to(tl.uint32, bitcast=True) + x_values_0 = _load_fp16x8(x_ptr + groups * GROUP_VALUES) + x_values_1 = _load_fp16x8(x_ptr + groups * GROUP_VALUES + 8) + x_values_2 = _load_fp16x8(x_ptr + groups * GROUP_VALUES + 16) + x_values_3 = _load_fp16x8(x_ptr + groups * GROUP_VALUES + 24) + partial = tl.zeros((GROUPS_PER_ROW,), tl.float32) + for value_index in tl.static_range(0, GROUP_VALUES): + bit_offset = value_index * 13 + word_index = bit_offset // 32 + shift = bit_offset % 32 + code = ( + _select_fp13_word( + word0, + word1, + word2, + word3, + word4, + word5, + word6, + word7, + word8, + word9, + word10, + word11, + word12, + INDEX=word_index, + ) + >> shift + ) + if shift > 19: + code |= _select_fp13_word( + word0, + word1, + word2, + word3, + word4, + word5, + word6, + word7, + word8, + word9, + word10, + word11, + word12, + INDEX=word_index + 1, + ) << (32 - shift) + code &= 0x1FFF + bits = (code << 3).to(tl.uint16) + weight = bits.to(tl.float16, bitcast=True).to(tl.float32) + if value_index < 8: + x = x_values_0[value_index] + elif value_index < 16: + x = x_values_1[value_index - 8] + elif value_index < 24: + x = x_values_2[value_index - 16] + else: + x = x_values_3[value_index - 24] + partial += x * weight + result = tl.sum(partial, axis=0) + tl.store(out_ptr + row, result) + + +@torch.no_grad() +def _pack_sm70_dsv4_fp13_weight(weight: torch.Tensor) -> torch.Tensor: + if weight.dtype != torch.float16 or weight.ndim != 2 or weight.shape[1] != _K: + raise ValueError( + f"unsupported FP13 weight contract: {weight.dtype=} {weight.shape=}" + ) + raw = weight.view(torch.int16).to(torch.int32) & 0xFFFF + codes = (raw >> 3).reshape( + weight.shape[0], _FP13_GROUPS_PER_ROW, _FP13_GROUP_VALUES + ) + words = torch.zeros( + (weight.shape[0], _FP13_GROUPS_PER_ROW, _FP13_GROUP_WORDS), + dtype=torch.int32, + device=weight.device, + ) + for value_index in range(_FP13_GROUP_VALUES): + bit_offset = value_index * 13 + word_index = bit_offset // 32 + shift = bit_offset % 32 + value = codes[..., value_index] + words[..., word_index].bitwise_or_(value << shift) + if shift > 19: + words[..., word_index + 1].bitwise_or_(value >> (32 - shift)) + return words.contiguous() + + +@torch.no_grad() +def _has_sm70_dsv4_fp13_weight_contract(weight: torch.Tensor) -> bool: + """Accept exact BF16-derived normals and bounded FP16 subnormal tails.""" + bits = weight.view(torch.int16).to(torch.int32) & 0xFFFF + exponent = (bits >> 10) & 0x1F + discarded = bits & 0x7 + incompatible = (exponent == 0x1F) | ((discarded != 0) & (exponent != 0)) + return not bool(torch.any(incompatible).item()) + + +def prepare_sm70_dsv4_fp13_gemv(layer: torch.nn.Module) -> torch.Tensor | None: + if not envs.VLLM_SM70_DSV4_FP13_GEMV or not getattr( + layer, "_sm70_dsv4_fp13_gemv", False + ): + return None + weight = getattr(layer, "weight", None) + if not isinstance(weight, torch.Tensor): + return None + if ( + not current_platform.is_cuda() + or not current_platform.is_device_capability((7, 0)) + or not weight.is_cuda + or weight.dtype != torch.float16 + or weight.ndim != 2 + or weight.shape[0] not in _SUPPORTED_N + or weight.shape[1] != _K + or not weight.is_contiguous() + ): + return None + if not _has_sm70_dsv4_fp13_weight_contract(weight): + logger.warning_once( + "SM70 packed-FP13 GEMV rejected a weight outside its tensor-value " + "contract; retaining the FP16 GEMV fallback." + ) + return None + packed = _pack_sm70_dsv4_fp13_weight(weight) + if _FP13_BUFFER in layer._buffers: + layer._buffers[_FP13_BUFFER] = packed + else: + layer.register_buffer(_FP13_BUFFER, packed, persistent=False) + return packed + + +def _has_sm70_dsv4_gemv_contract( x: torch.Tensor, weight: torch.Tensor, output_dtype: torch.dtype, ) -> bool: return ( - envs.VLLM_SM70_DSV4_FP16_GEMV - and current_platform.is_cuda() + current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) + and x.is_cuda + and weight.is_cuda + and x.device == weight.device and x.dtype == torch.float16 and weight.dtype == torch.float16 and output_dtype in (torch.float16, torch.float32) @@ -59,11 +289,57 @@ def can_use_sm70_dsv4_fp16_gemv( ) +def can_use_sm70_dsv4_fp16_gemv( + x: torch.Tensor, + weight: torch.Tensor, + output_dtype: torch.dtype, +) -> bool: + enabled = envs.VLLM_SM70_DSV4_FP16_GEMV or envs.VLLM_SM70_DSV4_FP13_GEMV + return enabled and _has_sm70_dsv4_gemv_contract(x, weight, output_dtype) + + +def can_use_sm70_dsv4_fp13_gemv( + x: torch.Tensor, + weight: torch.Tensor, + packed_weight: torch.Tensor | None, + output_dtype: torch.dtype, +) -> bool: + return ( + envs.VLLM_SM70_DSV4_FP13_GEMV + and _has_sm70_dsv4_gemv_contract(x, weight, output_dtype) + and packed_weight is not None + and packed_weight.dtype == torch.int32 + and packed_weight.device == weight.device + and packed_weight.shape + == (weight.shape[0], _FP13_GROUPS_PER_ROW, _FP13_GROUP_WORDS) + and packed_weight.is_contiguous() + ) + + def maybe_sm70_dsv4_fp16_gemv( x: torch.Tensor, weight: torch.Tensor, output_dtype: torch.dtype, + packed_weight: torch.Tensor | None = None, ) -> torch.Tensor | None: + if can_use_sm70_dsv4_fp13_gemv(x, weight, packed_weight, output_dtype): + assert packed_weight is not None + logger.info_once( + "DeepSeek V4 SM70 packed-FP13 GEMV enabled for batch-one decode." + ) + out = torch.empty((1, weight.shape[0]), device=x.device, dtype=output_dtype) + _sm70_dsv4_fp13_gemv_kernel[(weight.shape[0],)]( + x, + packed_weight, + out, + GROUP_VALUES=_FP13_GROUP_VALUES, + GROUP_WORDS=_FP13_GROUP_WORDS, + GROUPS_PER_ROW=_FP13_GROUPS_PER_ROW, + WORDS_PER_ROW=_FP13_WORDS_PER_ROW, + num_warps=_NUM_WARPS, + ) + return out + if not can_use_sm70_dsv4_fp16_gemv(x, weight, output_dtype): return None diff --git a/vllm/models/deepseek_v4/sm70/qnorm_rope_kv_fp8_insert.py b/vllm/models/deepseek_v4/sm70/qnorm_rope_kv_fp8_insert.py index 53f575d365..d520cf833f 100644 --- a/vllm/models/deepseek_v4/sm70/qnorm_rope_kv_fp8_insert.py +++ b/vllm/models/deepseek_v4/sm70/qnorm_rope_kv_fp8_insert.py @@ -232,7 +232,7 @@ def sm70_qnorm_rope_kv_fp8_insert( cos_sin_cache, cache_2d.stride(0), eps, - slot_mapping.shape[0], + num_tokens, num_heads=num_heads, cache_block_size=block_size, HEAD_DIM=_HEAD_DIM, diff --git a/vllm/models/deepseek_v4/sm70/sparse_kernels.py b/vllm/models/deepseek_v4/sm70/sparse_kernels.py index a32a94e118..1f741c88f1 100644 --- a/vllm/models/deepseek_v4/sm70/sparse_kernels.py +++ b/vllm/models/deepseek_v4/sm70/sparse_kernels.py @@ -38,6 +38,7 @@ def _sm70_sparse_gathered_kernel( BLOCK_H: tl.constexpr, BLOCK_K: tl.constexpr, BLOCK_D: tl.constexpr, + HAS_SINK: tl.constexpr, ): query_idx = tl.program_id(0) head_offsets = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) @@ -87,12 +88,16 @@ def _sm70_sparse_gathered_kernel( running_sum = running_sum * alpha + tl.sum(probs, axis=1) running_max = new_max - sink = tl.load(sink_ptr + head_offsets, mask=head_mask, other=neg_large).to( - tl.float32 - ) - final_max = tl.maximum(running_max, sink) - alpha = tl.exp(running_max - final_max) - final_sum = running_sum * alpha + tl.exp(sink - final_max) + if HAS_SINK: + sink = tl.load(sink_ptr + head_offsets, mask=head_mask, other=neg_large).to( + tl.float32 + ) + final_max = tl.maximum(running_max, sink) + alpha = tl.exp(running_max - final_max) + final_sum = running_sum * alpha + tl.exp(sink - final_max) + else: + alpha = tl.full((BLOCK_H,), 1.0, dtype=tl.float32) + final_sum = running_sum denom = tl.maximum(final_sum, 1.0e-30) result = tl.where( final_sum[:, None] > 0.0, @@ -836,7 +841,7 @@ def sm70_sparse_attention_gathered( indices: torch.Tensor, lengths: torch.Tensor, scale: float, - attn_sink: torch.Tensor, + attn_sink: torch.Tensor | None, out: torch.Tensor, ) -> None: """Sparse FP16 attention over a contiguous gathered KV workspace.""" @@ -848,12 +853,13 @@ def sm70_sparse_attention_gathered( assert indices_2d.shape[0] == q.shape[0] == lengths.shape[0] block_h = 8 + sink_arg = q if attn_sink is None else attn_sink.contiguous() _sm70_sparse_gathered_kernel[(q.shape[0], triton.cdiv(q.shape[1], block_h))]( q, kv_2d, indices_2d, lengths, - attn_sink.contiguous(), + sink_arg, out, q.stride(0), q.stride(1), @@ -868,6 +874,7 @@ def sm70_sparse_attention_gathered( BLOCK_H=block_h, BLOCK_K=16, BLOCK_D=_HEAD_DIM, + HAS_SINK=attn_sink is not None, num_warps=4, ) diff --git a/vllm/models/glm5next/__init__.py b/vllm/models/glm5next/__init__.py new file mode 100644 index 0000000000..dbd67b20f2 --- /dev/null +++ b/vllm/models/glm5next/__init__.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.platforms import current_platform + +if current_platform.is_rocm() or current_platform.is_xpu(): + raise NotImplementedError("GLM-5.3-Flash currently supports NVIDIA CUDA only.") + +from .nvidia.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration +from .nvidia.mtp import Glm5NextMTP + +__all__ = [ + "Glm5NextForCausalLM", + "Glm5NextForConditionalGeneration", + "Glm5NextMTP", +] diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/nvidia/attention.py new file mode 100644 index 0000000000..e28077f73d --- /dev/null +++ b/vllm/models/glm5next/nvidia/attention.py @@ -0,0 +1,619 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.config import ( + CacheConfig, + VllmConfig, +) +from vllm.distributed import ( + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope +from vllm.model_executor.layers.sparse_attn_indexer_kpool import SparseAttnIndexerKpool +from vllm.model_executor.models.deepseek_v2 import ( + DeepSeekV2FusedQkvAProjLinear, + DeepseekV32IndexerCache, + yarn_get_mscale, +) +from vllm.model_executor.models.utils import extract_layer_index +from vllm.model_executor.utils import maybe_disable_graph_partition +from vllm.models.glm5next.nvidia.ops.kpool_compress import ( + fwht128, + fwht128_quant_fp8, +) +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig +from vllm.v1.kv_cache_interface import KpoolTailSpec, MLAAttentionSpec + +logger = init_logger(__name__) + +# Shared torch.compile config for the indexer's small-kernel leaves. The MLA +# indexer runs under breakable-CG (CompilationMode.NONE), which blocks FX-graph +# fusion of the surrounding eager ops; carving each cluster into its own +# @torch.compile leaf (backend==inductor) still fuses them. Matches the +# grouped_topk / _cast_sigmoid leaf pattern. +_INDEXER_COMPILE = dict( + dynamic=True, + backend=current_platform.simple_compile_backend, + options=maybe_disable_graph_partition(current_platform.simple_compile_backend), +) + + +@torch.compile(**_INDEXER_COMPILE) +def _fused_indexer_k_norm( + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, dim: int, eps: float +) -> torch.Tensor: + # Fuse fp32 cast + layer_norm + cast-back (was 3 kernels) into one. + return F.layer_norm(x.float(), (dim,), weight, bias, eps).type_as(x) + + +@torch.compile(**_INDEXER_COMPILE) +def _fused_indexer_weight_scale( + weights: torch.Tensor, q_scale: torch.Tensor, scale: float +) -> torch.Tensor: + # Fuse the weight-scaling muls (was 2 kernels) into one. `scale` folds + # softmax_scale (head_dim**-0.5) and n_head**-0.5 into a single constant. + return (weights.unsqueeze(-1) * q_scale * scale).squeeze(-1) + + +@torch.compile(**_INDEXER_COMPILE) +def _fused_indexer_weight_only_scale( + weights: torch.Tensor, scale: float +) -> torch.Tensor: + return weights * scale + + +@torch.compile(**_INDEXER_COMPILE) +def _pad_indexer_heads(x: torch.Tensor, pad: int) -> torch.Tensor: + # DeepGEMM MQA-logits needs num_heads in {32,64}; zero-pad the head dim. + # Fuse new_zeros + cat (was 2 kernels) into one. Pad values are zero (exact + # in fp8 e4m3 and zero-weight in the logits sum), so numerically a no-op. + return torch.cat([x, x.new_zeros(x.shape[0], pad, *x.shape[2:])], dim=1) + + +class Glm5NextIndexerCache(DeepseekV32IndexerCache): + """Indexer K cache that stores kpool-compressed entries. + + Setting ``compress_ratio = index_kpool`` on the kv_cache_spec makes vLLM's + indexer metadata builder emit pool-granular ``slot_mapping`` / + ``seq_lens`` / ``cu_seq_lens`` / ``page_table`` for free, and shrinks the + cache allocation to ``storage_block_size = block_size // kpool``. The pool + *content* (softmax-weighted sum vs keep-every-Nth) is computed by the + kpool compress kernel inside the indexer op — the cache only provides the + addressing, which is identical for both schemes. + + The indexer shares one block with the co-located MLA (a single + ``MLAAttentionSpec`` / block_table), so ``block_size`` is the model-wide + ``cache_config.block_size``. DeepGEMM's paged-MQA kernel + (``csrc/apis/attention.hpp``) requires ``block_kv`` to be exactly 32 or + 64, so the storage block is virtually split into pool pages of the + largest such size that tiles it (``storage_kernel_block_size``); this + needs ``block_size`` to be a multiple of ``index_kpool * 32`` (512 for + ``index_kpool = 16``). A smaller block (e.g. the default 64) silently + collapses ``storage_block_size`` (64 // 16 = 4) and only fails later at + the opaque C++ assert; ``get_kv_cache_spec`` guards this up front + instead. + """ + + def __init__( + self, + *, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config, + index_kpool: int, + ): + super().__init__( + head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config + ) + assert index_kpool > 1, "Glm5NextIndexerCache expects index_kpool > 1" + # Keep chunked-prefill boundaries aligned to complete pools. + assert cache_config.block_size % index_kpool == 0, ( + "Glm5NextIndexerCache: cache_config.block_size " + f"({cache_config.block_size}) must be a multiple of index_kpool " + f"({index_kpool}) so chunked-prefill boundaries stay pool-aligned." + ) + self._index_kpool = index_kpool + + def get_kv_cache_spec(self, vllm_config: VllmConfig): + from dataclasses import replace + + spec = super().get_kv_cache_spec(vllm_config) + # compress_ratio lives on MLAAttentionSpec, but the base + # DeepseekV32IndexerCache.get_kv_cache_spec is typed to return the + # KVCacheSpec base; narrow so dataclass.replace sees the field. + assert isinstance(spec, MLAAttentionSpec) + spec = replace(spec, compress_ratio=self._index_kpool) + + # DeepGEMM paged-MQA takes block_kv in {32, 64}; the storage block + # (= block_size // index_kpool) is virtually split into pool pages of + # the largest such size that tiles it, so it must be a multiple of 32. + storage_block_size = spec.block_size // self._index_kpool + assert ( + spec.block_size % self._index_kpool == 0 and storage_block_size % 32 == 0 + ), ( + "Glm5NextIndexerCache: kpool indexer requires cache block_size to " + f"be a multiple of index_kpool * 32 ({self._index_kpool * 32}) so " + "that DeepGEMM paged-MQA pool pages (32 or 64 entries) tile the " + f"storage block, got block_size={spec.block_size} -> " + f"storage_block_size={storage_block_size}." + ) + return spec + + +class Glm5NextTailCache(DeepseekV32IndexerCache): + """Paged circular buffer for the kpool indexer's in-progress (tail) pool. + + Holds the trailing incomplete pool's raw K + gate score: one block of + ``index_kpool`` slots per request, overwritten in place by ``pos % kpool`` + as decode/spec-decode advances. Prefill seeds it (instead of discarding the + tail raw K+gate); the connector transfers it across PD; decode reads it to + compress the boundary pool correctly. ``KpoolTailSpec`` / + ``KpoolTailManager`` provide the no-prune, 1-block/req allocation that lets + the in-progress pool survive across steps and across transfer. + + Stores raw model-dtype K (``head_dim``) as the "K" half of each block and + gate score (``head_dim``) as the "V" half -- not the fp8-compressed + entry, which lives in ``Glm5NextIndexerCache``. + """ + + def __init__( + self, + *, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config, + index_kpool: int, + ): + super().__init__( + head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config + ) + assert index_kpool > 1, "Glm5NextTailCache expects index_kpool > 1" + self._index_kpool = index_kpool + + def get_kv_cache_spec(self, vllm_config: VllmConfig): + # K + gate score packed into head_size (== 2*head_dim), head_size_v=0: + # KpoolTailBackend.get_kv_cache_shape only consumes head_size and splits + # it into [2, kpool, head_dim] (K | score halves), so the connectors' + # non-MLA K/V half-split transfers K and score as separate halves. + return KpoolTailSpec( + block_size=self._index_kpool, + num_kv_heads=1, + head_size=2 * self.head_dim, + head_size_v=0, + dtype=self.dtype, + sliding_window=self._index_kpool, + ) + + def get_attn_backend(self): + from vllm.v1.attention.backends.mla.indexer import KpoolTailBackend + + return KpoolTailBackend + + +class Indexer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: Glm5NextConfig, + hidden_size: int, + q_lora_rank: int, + quant_config: QuantizationConfig | None, + cache_config: CacheConfig | None, + topk_indices_buffer: torch.Tensor | None, + prefix: str = "", + ): + super().__init__() + self.vllm_config = vllm_config + self.config = config + self.quant_config = quant_config + # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"] + # Indexer is only constructed for v32 configs, where these sparse-indexer + # fields are guaranteed populated; narrow away the `int | None` declared + # on Glm5NextConfig for the optional-indexer case. + assert config.index_topk is not None + assert config.index_n_heads is not None + assert config.index_head_dim is not None + assert config.index_kpool is not None + self.topk_tokens = config.index_topk + self.n_head = config.index_n_heads # 64 + self.head_dim = config.index_head_dim # 128 + self.rope_dim = config.qk_rope_head_dim # 64 + self.index_kpool = config.index_kpool + self.q_lora_rank = q_lora_rank # 1536 + self.activation_dtype = vllm_config.model_config.dtype + + # kpool + self.index_kpool_compress_ape = nn.Parameter( + torch.zeros(self.index_kpool, self.head_dim, dtype=torch.float32) + ) + # Keep the checkpoint name ``index_kpool_compress_gate`` without a + # ``.weight`` suffix. F.linear consumes its [head_dim, hidden_size] shape. + self.index_kpool_compress_gate = nn.Parameter( + torch.empty(self.head_dim, hidden_size, dtype=self.activation_dtype) + ) + + # no tensor parallel, just replicated + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.head_dim * self.n_head, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wq_b", + ) + # Fused wk + weights_proj: single GEMM producing [head_dim + n_head]. + # FP8 wk weights are upcasted to BF16 during loading to maintain fusion. + self.wk_weights_proj = MergedColumnParallelLinear( + hidden_size, + [self.head_dim, self.n_head], + bias=False, + quant_config=None, + disable_tp=True, + prefix=f"{prefix}.wk_weights_proj", + ) + self.k_norm = LayerNorm(self.head_dim, eps=1e-6) + self.softmax_scale = self.head_dim**-0.5 + + # Hadamard-128 rotation of the indexer query is fused with the FP8 + # quant (see forward: fwht128_quant_fp8) -- no precomputed matrix. + + self.scale_fmt = "ue8m0" + self.quant_block_size = 128 # TODO: get from config + self.topk_indices_buffer = topk_indices_buffer + + # NOTE: (zyongye) we use fp8 naive cache, + # where we store value in fp8 and scale in fp32 + # per self.quant_block_size element + self.k_cache = Glm5NextIndexerCache( + head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4, + dtype=torch.uint8, + prefix=f"{prefix}.k_cache", + cache_config=cache_config, + index_kpool=self.index_kpool, + ) + # Paged tail cache (in-progress pool's raw K + gate score). Written by + # prefill (seeds the boundary pool) and decode (per-step stash); read by + # the decode kernel to compress the boundary pool. Transferred across PD + # so the decode side sees the prefill tail. See KpoolTailSpec/Manager. + self.tail_cache = Glm5NextTailCache( + head_dim=self.head_dim, + dtype=self.activation_dtype, + prefix=f"{prefix}.tail_cache", + cache_config=cache_config, + index_kpool=self.index_kpool, + ) + self.max_model_len = vllm_config.model_config.max_model_len + self.prefix = prefix + from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size + + self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config) + self.indexer_op = SparseAttnIndexerKpool( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + tail_cache=self.tail_cache, + ) + + def forward( + self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb + ) -> torch.Tensor: + q, _ = self.wq_b(qr) + q = q.view(-1, self.n_head, self.head_dim) + + # Compute the head gate in fp32; bf16 error can change near-tie pool + # rankings on long-context tasks. Cache it after weights are loaded. + kw, _ = self.wk_weights_proj(hidden_states) + k = kw[:, : self.head_dim] + if getattr(self, "_wp_fp32", None) is None: + self._wp_fp32 = ( + self.wk_weights_proj.weight.data[self.head_dim :, :] + .t() + .contiguous() + .float() + ) + weights = torch.mm(hidden_states.float(), self._wp_fp32) + + k = _fused_indexer_k_norm( + k, self.k_norm.weight, self.k_norm.bias, self.head_dim, self.k_norm.eps + ) + + if self.rope_dim > 0: + q_pe, q_nope = torch.split( + q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + k_pe, k_nope = torch.split( + k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1)) + # Note: RoPE (NeoX) can introduce extra leading dimensions during + # compilation so we need to reshape back to token-flattened shapes + q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim) + k_pe = k_pe.reshape(-1, 1, self.rope_dim) + + # `rotary_emb` is shape-preserving; `q_pe` is already + # [num_tokens, n_head, rope_dim]. + q = torch.cat([q_pe, q_nope], dim=-1) + # `k_pe` is [num_tokens, 1, rope_dim] (MQA). + k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) + # else: qk_rope_head_dim=0 — no rope component. q is already + # [num_tokens, n_head, head_dim] and k is [num_tokens, head_dim] (all + # nope), so skip the rope split / rotary / cat entirely; otherwise the + # split/reshape would build 0-element tensors (breaks dynamo tracing). + + # Rotate Q into the cached K basis before computing fp8 MQA logits. + # Fusing the fp32 FWHT and quantization avoids an intermediate HBM + # round-trip and bf16 matrix-rounding bias. + assert self.head_dim == 128 and self.quant_block_size == 128 + assert self.scale_fmt == "ue8m0" + q = q.view(-1, self.head_dim) + sm70_fp16_indexer = current_platform.is_cuda() and ( + current_platform.is_device_capability((7, 0)) + ) + if sm70_fp16_indexer: + assert q.dtype == torch.float16 + q_index = fwht128(q).view(-1, self.n_head, self.head_dim) + weights = _fused_indexer_weight_only_scale( + weights, self.softmax_scale * self.n_head**-0.5 + ) + else: + q_fp8, q_scale = fwht128_quant_fp8(q) + q_index = q_fp8.view(-1, self.n_head, self.head_dim) + q_scale = q_scale.view(-1, self.n_head, 1) + weights = _fused_indexer_weight_scale( + weights, q_scale, self.softmax_scale * self.n_head**-0.5 + ) + + # kpool: per-token gate score driving the softmax-weighted pool. Computed + # from the same hidden_states that produced `k`, so it stays token-aligned. + # F.linear(x, gate) = x @ gate.T with gate [head_dim, hidden_size]. + gate_score = F.linear(hidden_states, self.index_kpool_compress_gate) + + # DeepGEMM's MQA-logits kernels (fp8_mqa_logits / + # fp8_fp4_paged_mqa_logits) require num_heads in {32, 64}; this + # checkpoint uses index_n_heads=16. Zero-pad q and the per-head + # weights: logits are a weights-weighted sum over heads, so + # zero-weight padded heads contribute exactly nothing. + if self.n_head < 32: + pad = 32 - self.n_head + q_index = _pad_indexer_heads(q_index, pad) + weights = _pad_indexer_heads(weights, pad) + + return self.indexer_op( + hidden_states, + q_index, + k, + weights, + gate_score=gate_score, + compress_ape=self.index_kpool_compress_ape, + index_kpool=self.index_kpool, + positions=positions, + ) + + +class Glm5NextMLAAttention(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: Glm5NextConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + max_position_embeddings: int = 8192, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + topk_indices_buffer: torch.Tensor | None = None, + input_size: int | None = None, + skip_rope: bool | None = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + + self.num_heads = num_heads + tp_size = get_tensor_model_parallel_world_size() + assert num_heads % tp_size == 0 + self.num_local_heads = num_heads // tp_size + + self.scaling = self.qk_head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + + # Use input_size for projection input dimensions if provided, + # otherwise default to hidden_size (used in Eagle3 Deepseek with MLA) + proj_input_size = input_size if input_size is not None else self.hidden_size + + if self.q_lora_rank is not None: + self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear( + proj_input_size, + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + ) + else: + self.kv_a_proj_with_mqa = ReplicatedLinear( + proj_input_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + + if self.q_lora_rank is not None: + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + self.q_proj = ColumnParallelLinear( + proj_input_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + if not skip_rope: + assert config.rope_parameters is not None + if config.rope_parameters["rope_type"] != "default": + config.rope_parameters["rope_type"] = ( + "deepseek_yarn" + if config.rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + + self.rotary_emb: RotaryEmbedding | None = get_rope( + qk_rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=False, + ) + + if ( + config.rope_parameters["rope_type"] != "default" + and config.rope_parameters["rope_type"] == "deepseek_yarn" + ): + mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False) + scaling_factor = config.rope_parameters["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale + else: + self.rotary_emb = None + + # `index_topk` is declared on Glm5NextTextConfig with a default of None, + # so hasattr() is True even for full-MLA configs (no kpool indexer). + self.is_v32 = getattr(config, "index_topk", None) is not None + # self.is_v32 = False + + _skip_topk = False + if self.is_v32: + self.indexer_rope_emb: RotaryEmbedding | None = get_rope( + qk_rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=not getattr(config, "indexer_rope_interleave", False), + ) + # The sparse indexer projects from the MLA q-lora rank, which is + # always set for v32 MLA configs; narrow away the `int | None`. + assert q_lora_rank is not None + self.indexer: Indexer | None = Indexer( + vllm_config, + config, + hidden_size, + q_lora_rank, + quant_config, + cache_config, + topk_indices_buffer, + f"{prefix}.indexer", + ) + + # Enable IndexCache for DeepSeek models to reduce redundant top-k + # token selection computations in sparse attention. + use_index_cache = getattr(config, "use_index_cache", False) + if use_index_cache: + # IndexCache config + # Refer: https://arxiv.org/abs/2603.12201 for more details. + _index_topk_freq = getattr(config, "index_topk_freq", 1) + _index_topk_pattern = getattr(config, "index_topk_pattern", None) + layer_id = extract_layer_index(prefix) + if _index_topk_pattern is None: + _skip_topk = max(layer_id - 1, 0) % _index_topk_freq != 0 + elif 0 <= layer_id < len(_index_topk_pattern): + _skip_topk = _index_topk_pattern[layer_id] == "S" + + else: + self.indexer_rope_emb = None + self.indexer = None + + mla_modules = MLAModules( + kv_a_layernorm=self.kv_a_layernorm, + kv_b_proj=self.kv_b_proj, + rotary_emb=self.rotary_emb, + o_proj=self.o_proj, + fused_qkv_a_proj=self.fused_qkv_a_proj + if self.q_lora_rank is not None + else None, + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa + if self.q_lora_rank is None + else None, + q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None, + q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None, + q_proj=self.q_proj if self.q_lora_rank is None else None, + indexer=self.indexer, + indexer_rotary_emb=self.indexer_rope_emb, + is_sparse=self.is_v32, + topk_indices_buffer=topk_indices_buffer, + ) + + self.mla_attn = MultiHeadLatentAttentionWrapper( + self.hidden_size, + self.num_local_heads, + self.scaling, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + self.q_lora_rank, + self.kv_lora_rank, + mla_modules, + cache_config, + quant_config, + prefix, + skip_topk=_skip_topk, + fuse_qkv_rmsnorm=True, + ) + + def forward( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + # The wrapper also runs the sparse indexer before MLA attention. + return self.mla_attn(positions, hidden_states) diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py new file mode 100644 index 0000000000..22732a5015 --- /dev/null +++ b/vllm/models/glm5next/nvidia/kda.py @@ -0,0 +1,675 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-5.3-Flash KDA layer with separate convolutions and a bounded safe gate.""" + +import torch +from torch import nn + +import vllm._sm70_ops as sm70_ops +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import divide +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.fla.ops.kda import ( + FusedRMSNormGated, + chunk_kda_with_fused_gate, + fused_recurrent_kda, +) +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateDtypeCalculator, + MambaStateShapeCalculator, + is_conv_state_dim_first, +) +from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( + gather_initial_states, +) +from vllm.model_executor.layers.mamba.ops.scatter_states import scatter_states +from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader +from vllm.model_executor.utils import ( + maybe_disable_graph_partition, + set_weight_attrs, +) +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata + +logger = init_logger(__name__) + + +class _Glm5NextMergedColumnParallelLinear(MergedColumnParallelLinear): + """Merged projection with multiple replicated output shards. + + Extends K3's ``_KimiGDNMergedColumnParallelLinear`` to support two + replicated shards (f_a, g_a) instead of one. Pre-multiplies each + replicated entry's output_size by tp_size so the per-rank shard + divides back to the full size, and forces tp_rank=0 during weight + loading for replicated shards. + """ + + def __init__( + self, + input_size: int, + output_sizes: list[int], + replicated_shard_ids: tuple[int, ...], + tp_size: int, + **kwargs, + ) -> None: + self.replicated_shard_ids = set(replicated_shard_ids) + output_sizes = output_sizes.copy() + for sid in self.replicated_shard_ids: + output_sizes[sid] *= tp_size + super().__init__(input_size, output_sizes, **kwargs) + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id in self.replicated_shard_ids: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + def weight_loader_v2( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id in self.replicated_shard_ids: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader_v2(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + +@torch.compile( + dynamic=True, + backend=current_platform.simple_compile_backend, + options=maybe_disable_graph_partition(current_platform.simple_compile_backend), +) +def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor: + """Fuse the fp32 cast + sigmoid into one Inductor kernel.""" + return x.float().sigmoid() + + +class Glm5NextLinearAttention(GatedDeltaNetAttention): + # Declared int (set in __init__ from config) so mypy doesn't see the + # getattr-derived `Any | None` at the kernel call sites. + head_dim: int + num_heads: int + conv_size: int + + def get_state_dtype( + self, + ) -> tuple[torch.dtype, torch.dtype]: + if self.model_config is None or self.cache_config is None: + raise ValueError("model_config and cache_config must be set") + q_dtype, _, _, recurrent_dtype = MambaStateDtypeCalculator.kda_state_dtype( + self.model_config.dtype, self.cache_config.mamba_cache_dtype + ) + return q_dtype, recurrent_dtype + + def get_state_shape( + self, + ) -> tuple[tuple[int, ...], tuple[int, ...]]: + # conv_state width must include num_spec so the spec-decode conv update + # (causal_conv1d_update with num_accepted_tokens + max_query_len) can + # slide the window across the draft-verify tokens without reading past + # the allocated width. Matches qwen_gdn_linear_attn.get_state_shape. + q_shape, k_shape, v_shape, recurrent_shape = ( + MambaStateShapeCalculator.kda_state_shape( + self.tp_size, + self.num_heads, + self.head_dim, + conv_kernel_size=self.conv_size, + num_spec=self.num_spec, + ) + ) + if is_conv_state_dim_first(): + merged_conv_shape = ( + q_shape[0] + k_shape[0] + v_shape[0], + q_shape[1], + ) + else: + merged_conv_shape = ( + q_shape[0], + q_shape[1] + k_shape[1] + v_shape[1], + ) + return merged_conv_shape, recurrent_shape + + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + # KDA projections remain BF16 because fp8 checkpoints omit their scales. + saved_quant_config = vllm_config.quant_config + vllm_config.quant_config = None + super().__init__(config, vllm_config, prefix) + vllm_config.quant_config = saved_quant_config + + # Linear-attention head config: read the flattened top-level fields when + # present (new schema); fall back to the legacy linear_attn_config dict + # otherwise (shared base is also used by KimiLinearConfig). Narrow via + # locals so the int-typed attrs are assigned a non-None value. + head_dim = getattr(config, "linear_head_dim", None) + num_heads = getattr(config, "linear_num_heads", None) + conv_size = getattr(config, "linear_conv_kernel_dim", None) + if head_dim is None or num_heads is None or conv_size is None: + kda_config = config.linear_attn_config # type: ignore[attr-defined] + assert kda_config is not None, "linear_attn_config must be set" + head_dim = kda_config["head_dim"] + num_heads = kda_config["num_heads"] + conv_size = kda_config["short_conv_kernel_size"] + assert head_dim is not None + assert num_heads is not None + assert conv_size is not None + self.head_dim = head_dim + self.num_heads = num_heads + self.conv_size = conv_size + assert self.num_heads % self.tp_size == 0 + self.local_num_heads = divide(self.num_heads, self.tp_size) + + projection_size = self.head_dim * self.num_heads + self.local_projection_size = divide(projection_size, self.tp_size) + self._use_sm70_fused_fg_b_decode = ( + current_platform.is_cuda() + and current_platform.get_device_capability() == (7, 0) + and self.tp_size == 4 + and self.head_dim == 128 + and self.local_projection_size == 2048 + ) + + # Merge q, k, v, b, f_a, g_a projections into one GEMM (6→1 launches). + # Order matches checkpoint's fused_qkvbfg_a_proj convention. + # Shards 4 (f_a) and 5 (g_a) are replicated across TP ranks. + self.in_proj_qkvbfg_a = _Glm5NextMergedColumnParallelLinear( + self.hidden_size, + [ + projection_size, # q (shard 0) + projection_size, # k (shard 1) + projection_size, # v (shard 2) + self.num_heads, # b (shard 3) + self.head_dim, # f_a (shard 4, replicated) + self.head_dim, # g_a (shard 5, replicated) + ], + replicated_shard_ids=(4, 5), + tp_size=self.tp_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.in_proj_qkvbfg_a", + ) + + self.f_b_proj = ColumnParallelLinear( + self.head_dim, + projection_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.f_b_proj", + ) + self.dt_bias = nn.Parameter( + torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32) + ) + + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + + self.q_conv1d = ColumnParallelLinear( + input_size=self.conv_size, + output_size=projection_size, + bias=False, + params_dtype=torch.float32, + prefix=f"{prefix}.q_conv1d", + ) + self.k_conv1d = ColumnParallelLinear( + input_size=self.conv_size, + output_size=projection_size, + bias=False, + params_dtype=torch.float32, + prefix=f"{prefix}.k_conv1d", + ) + self.v_conv1d = ColumnParallelLinear( + input_size=self.conv_size, + output_size=projection_size, + bias=False, + params_dtype=torch.float32, + prefix=f"{prefix}.v_conv1d", + ) + # unsqueeze to fit conv1d weights shape into the linear weights shape. + # Can't do this in `weight_loader` since it already exists in + # `ColumnParallelLinear` and `set_weight_attrs` + # doesn't allow to override it + self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1) + self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1) + self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1) + # Lazily-built merged q|k|v conv weight (built on first forward, after + # weights are loaded). See _forward. + self._merged_conv_weight: torch.Tensor | None = None + + self.A_log = nn.Parameter( + torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32) + ) + set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)}) + + self.g_b_proj = ColumnParallelLinear( + self.head_dim, + projection_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.g_b_proj", + ) + self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") + self.o_proj = RowParallelLinear( + projection_size, + self.hidden_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.o_proj", + ) + + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + # Checkpoints store A_log as 1-D; the model parameter is 4-D. + def _a_log_weight_loader(param, loaded_weight): + if loaded_weight.dim() == 1: + loaded_weight = loaded_weight.view([1, 1, -1, 1]) + return sharded_weight_loader(2)(param, loaded_weight) + + self.A_log.weight_loader = _a_log_weight_loader + + # GLM-5.3-Flash uses a bounded sigmoid gate instead of the default + # unbounded softplus gate. + linear_lower_bound = getattr(config, "linear_lower_bound", None) + if linear_lower_bound is not None: + self.kda_safe_gate = True + self.kda_lower_bound = linear_lower_bound + else: + legacy = getattr(config, "linear_attn_config", None) or {} + if legacy.get("safe_gate", True): + self.kda_safe_gate = True + self.kda_lower_bound = legacy.get("lower_bound", -5.0) + else: + self.kda_safe_gate = False + self.kda_lower_bound = -5.0 + # Process-global conv-state layout, resolved once here instead of on + # every _forward call (it reads an env-derived flag each time). + self._conv_state_dim_first = is_conv_state_dim_first() + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + num_tokens = hidden_states.size(0) + # One merged GEMM for q, k, v, b, f_a, g_a (replaces 6 separate GEMMs). + projected = self.in_proj_qkvbfg_a(hidden_states)[0] + qkv, beta_raw, f_a, g_a = projected.split( + [ + 3 * self.local_projection_size, + self.local_num_heads, + self.head_dim, + self.head_dim, + ], + dim=-1, + ) + + # Beta stays raw (bf16) here: the recurrent kernel sigmoids it in fp32 + # at load (SIGMOID_BETA), and only the chunked prefill path needs the + # pre-computed fp32 sigmoid — computed lazily in _forward. Pure decode + # / spec-verify steps then skip the _cast_sigmoid kernel and its fp32 + # intermediate entirely. + beta = beta_raw.unsqueeze(0) + if self._use_sm70_fused_fg_b_decode and num_tokens == 1: + if not hasattr(torch.ops._C, "sm70_glm_kda_fg_b_out"): + raise RuntimeError( + "SM70 GLM KDA decode requires the native fused f_b/g_b op. " + "Rebuild vLLM from source with CUDA arch 7.0." + ) + g1 = torch.empty( + (1, self.local_projection_size), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + g_proj_states = torch.empty_like(g1) + sm70_ops.sm70_glm_kda_fg_b_out( + g1, + g_proj_states, + f_a, + g_a, + self.f_b_proj.weight, + self.g_b_proj.weight, + ) + logger.info_once("SM70 GLM KDA fused f_b/g_b decode path enabled.") + else: + g1 = self.f_b_proj(f_a)[0] + g_proj_states = self.g_b_proj(g_a)[0] + g1 = g1.reshape(1, -1, self.local_num_heads, self.head_dim) + + # Must stay 3D: rms_norm_gated reads H from g.shape[-2]. + g2 = g_proj_states.reshape(-1, self.local_num_heads, self.head_dim) + + core_attn_out = torch.empty( + (1, num_tokens, self.local_num_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + # Call the decorated eager break directly so host-side prefill branches + # are not captured by PIECEWISE CUDA graphs. + self._forward( + qkv_proj_states=qkv, + g1=g1, + beta=beta, + core_attn_out=core_attn_out, + ) + core_attn_out = self.o_norm(core_attn_out, g2) + core_attn_out = core_attn_out.reshape(core_attn_out.size(1), -1) + return self.o_proj(core_attn_out)[0] + + @eager_break_during_capture + def _forward( + self, + qkv_proj_states: torch.Tensor, + g1: torch.Tensor, + beta: torch.Tensor, + core_attn_out: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + + if attn_metadata_raw is None: + # # V1 profile run + return + + assert isinstance(attn_metadata_raw, dict) + attn_metadata_narrowed = attn_metadata_raw[self.prefix] + assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata) + has_initial_state = attn_metadata_narrowed.has_initial_state + non_spec_query_start_loc = attn_metadata_narrowed.non_spec_query_start_loc + non_spec_state_indices_tensor = ( + attn_metadata_narrowed.non_spec_state_indices_tensor + ) # noqa: E501 + num_actual_tokens = attn_metadata_narrowed.num_actual_tokens + # Spec-decode metadata (all None when speculative decoding is disabled). + spec_sequence_masks = attn_metadata_narrowed.spec_sequence_masks + spec_query_start_loc = attn_metadata_narrowed.spec_query_start_loc + spec_state_indices_tensor = attn_metadata_narrowed.spec_state_indices_tensor + spec_token_indx = attn_metadata_narrowed.spec_token_indx + non_spec_token_indx = attn_metadata_narrowed.non_spec_token_indx + num_accepted_tokens = attn_metadata_narrowed.num_accepted_tokens + num_spec_decodes = attn_metadata_narrowed.num_spec_decodes + use_spec = spec_sequence_masks is not None and num_spec_decodes > 0 + # Safe-gate checkpoints use the bounded sigmoid variant. + safe_gate = self.kda_safe_gate + lower_bound = self.kda_lower_bound + constant_caches = self.kv_cache + + qkv_proj_states = qkv_proj_states[:num_actual_tokens] + g1 = g1[:, :num_actual_tokens] + beta = beta[:, :num_actual_tokens] + + (conv_state, recurrent_state) = constant_caches + # conv_state must be (..., dim, width-1) for the conv kernels. + # DS layout stores it that way directly; SD layout needs a transpose. + # Layout is process-global and resolved once at init (see __init__). + if not self._conv_state_dim_first: + conv_state = conv_state.transpose(-1, -2) + + # One merged short-conv over q|k|v instead of three separate calls. The + # 1D conv is independent per channel, so concatenating q/k/v along the + # channel dim and running a single causal_conv1d is bit-identical to + # three calls. The merged weight is q|k|v conv weights concatenated; + # built once and cached (params are fixed after load). conv_state is + # already stored as the merged q|k|v state, so it is used directly. + if self._merged_conv_weight is None: + + def _w(m): + return m.weight.view(m.weight.size(0), m.weight.size(2)) + + self._merged_conv_weight = torch.cat( + [_w(self.q_conv1d), _w(self.k_conv1d), _w(self.v_conv1d)], + dim=0, + ).contiguous() + conv_weights = self._merged_conv_weight + conv_bias = self.q_conv1d.bias + + # Split projections / gating into spec (draft-verify) and non-spec token + # groups when speculative decoding is active. Spec tokens carry + # num_spec+1 recurrent-state columns each and are advanced with + # num_accepted_tokens for rejection-sampling rollback; non-spec tokens + # are one-per-request. Mirrors olmo_gdn_linear_attn.py. Projections are + # [n, *] (token dim 0); g1/beta are [1, n, h, d] (token dim 1). + if use_spec: + # In a pure spec-verify step (no non-spec tokens) the metadata + # builder sets spec_token_indx = arange(num_actual_tokens), making + # the index_select calls below identity copies. Skip them on this + # steady-state decode hot path. The outputs alias the inputs here; + # the downstream conv/recurrent kernels read them without mutating + # in place, so the aliasing is safe. + if non_spec_token_indx is None or non_spec_token_indx.numel() == 0: + qkv_spec = qkv_proj_states + g1_spec = g1 + beta_spec = beta + else: + qkv_spec = qkv_proj_states.index_select(0, spec_token_indx) + g1_spec = g1.index_select(1, spec_token_indx) + beta_spec = beta.index_select(1, spec_token_indx) + if non_spec_token_indx is not None and non_spec_token_indx.numel() > 0: + qkv_ns = qkv_proj_states.index_select(0, non_spec_token_indx) + g1_ns = g1.index_select(1, non_spec_token_indx) + beta_ns = beta.index_select(1, non_spec_token_indx) + else: + qkv_ns = g1_ns = beta_ns = None + else: + qkv_spec = g1_spec = beta_spec = None + qkv_ns, g1_ns, beta_ns = qkv_proj_states, g1, beta + + # --- causal conv1d: spec (draft-verify) path --- + if use_spec: + assert spec_state_indices_tensor is not None + assert num_accepted_tokens is not None + conv_idx = spec_state_indices_tensor[:, 0][:num_spec_decodes] + conv_mql = spec_state_indices_tensor.size(-1) + qkv_spec = causal_conv1d_update( + qkv_spec, + conv_state, + conv_weights, + conv_bias, + activation="silu", + conv_state_indices=conv_idx, + num_accepted_tokens=num_accepted_tokens, + query_start_loc=spec_query_start_loc, + max_query_len=conv_mql, + ) + q_spec, k_spec, v_spec = qkv_spec.split(self.local_projection_size, dim=-1) + + # --- causal conv1d: non-spec path (prefill or plain decode) --- + q_ns = k_ns = v_ns = None + if attn_metadata_narrowed.num_prefills > 0: + assert qkv_ns is not None + qkv_ns = causal_conv1d_fn( + qkv_ns.transpose(0, 1), + conv_weights, + conv_bias, + activation="silu", + conv_states=conv_state, + has_initial_state=has_initial_state, + cache_indices=non_spec_state_indices_tensor, + query_start_loc=non_spec_query_start_loc, + metadata=attn_metadata_narrowed, + ).transpose(0, 1) + q_ns, k_ns, v_ns = qkv_ns.split(self.local_projection_size, dim=-1) + elif attn_metadata_narrowed.num_decodes > 0: + assert non_spec_state_indices_tensor is not None + decode_conv_indices = non_spec_state_indices_tensor[ + : attn_metadata_narrowed.num_decodes + ] + qkv_ns = causal_conv1d_update( + qkv_ns, + conv_state, + conv_weights, + conv_bias, + activation="silu", + conv_state_indices=decode_conv_indices, + ) + q_ns, k_ns, v_ns = qkv_ns.split(self.local_projection_size, dim=-1) + + def _rearr(x): + return x.reshape(1, -1, self.local_num_heads, self.head_dim) + + # --- core attention: spec (draft-verify) path --- + core_attn_out_spec = None + # In a pure spec-verify step (no non-spec tokens) the recurrent kernel + # can write straight into the layer output buffer, skipping the + # fresh allocation + copy below. Mixed steps must scatter via + # spec_token_indx, so they keep the kernel-managed output. + spec_out = ( + core_attn_out[0, :num_actual_tokens].unsqueeze(0) + if non_spec_token_indx is None or non_spec_token_indx.numel() == 0 + else None + ) + if use_spec: + assert spec_state_indices_tensor is not None + assert num_accepted_tokens is not None + assert spec_query_start_loc is not None + # Gate computed inside the recurrent kernel (COMPUTE_GATE) from + # raw g1 — replicates fused_kda_gate's arithmetic bit-for-bit and + # skips its launch + fp32 [n, H, D] intermediate per layer. + core_attn_out_spec, _ = fused_recurrent_kda( + q=_rearr(q_spec), + k=_rearr(k_spec), + v=_rearr(v_spec), + g=g1_spec, + beta=beta_spec, + initial_state=recurrent_state, + use_qk_l2norm_in_kernel=True, + cu_seqlens=spec_query_start_loc[: num_spec_decodes + 1], + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + out=spec_out, + sigmoid_beta=True, + a_log=self.A_log, + g_bias=self.dt_bias, + compute_gate=True, + lower_bound=lower_bound, + ) + + # --- core attention: non-spec path (prefill or plain decode) --- + core_attn_out_non_spec = None + # Only the plain-decode recurrent kernel can write straight into the + # layer output buffer; the chunked prefill kernel cannot, so this + # stays None there and the merge copy below runs as before. + ns_out = None + if attn_metadata_narrowed.num_prefills > 0: + assert q_ns is not None + assert non_spec_state_indices_tensor is not None + assert has_initial_state is not None + initial_state = gather_initial_states( + recurrent_state, non_spec_state_indices_tensor, has_initial_state + ) + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = chunk_kda_with_fused_gate( + q=_rearr(q_ns), + k=_rearr(k_ns), + v=_rearr(v_ns), + raw_g=g1_ns, + # Chunk path wants the pre-sigmoided fp32 beta (its kernels + # don't sigmoid); beta_ns is raw bf16 from forward. + beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0), + A_log=self.A_log, + g_bias=self.dt_bias, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=non_spec_query_start_loc, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) + # Init cache + scatter_states( + recurrent_state, + last_recurrent_state, + non_spec_state_indices_tensor, + ) + elif attn_metadata_narrowed.num_decodes > 0: + assert non_spec_query_start_loc is not None + assert non_spec_state_indices_tensor is not None + # Plain decode step (no spec tokens): token order is dense, so the + # kernel can write straight into the layer output buffer. A mixed + # step scatters non-spec output via non_spec_token_indx instead. + # Gate computed in-kernel (COMPUTE_GATE), beta sigmoided in-kernel. + if not use_spec: + ns_out = spec_out + core_attn_out_non_spec, _ = fused_recurrent_kda( + q=_rearr(q_ns), + k=_rearr(k_ns), + v=_rearr(v_ns), + g=g1_ns, + beta=beta_ns, + initial_state=recurrent_state, + use_qk_l2norm_in_kernel=True, + cu_seqlens=non_spec_query_start_loc[ + : attn_metadata_narrowed.num_decodes + 1 + ], + ssm_state_indices=non_spec_state_indices_tensor, + out=ns_out, + sigmoid_beta=True, + a_log=self.A_log, + g_bias=self.dt_bias, + compute_gate=True, + lower_bound=lower_bound, + ) + + # --- merge spec / non-spec outputs back into token order --- + if use_spec and core_attn_out_non_spec is not None: + assert core_attn_out_spec is not None + merged = torch.empty( + (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), + dtype=core_attn_out_non_spec.dtype, + device=core_attn_out_non_spec.device, + ) + merged.index_copy_(1, spec_token_indx, core_attn_out_spec) + merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) + core_attn_out[0, :num_actual_tokens] = merged.squeeze(0) + elif use_spec: + assert core_attn_out_spec is not None + if spec_out is None: + core_attn_out[0, :num_actual_tokens] = core_attn_out_spec.squeeze(0) + else: + assert core_attn_out_non_spec is not None + if ns_out is None: + core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ + 0, :num_actual_tokens + ] diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py new file mode 100644 index 0000000000..ad08e80424 --- /dev/null +++ b/vllm/models/glm5next/nvidia/model.py @@ -0,0 +1,1286 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from typing import ClassVar, Literal + +import torch +from torch import nn + +from vllm.config import ParallelConfig, VllmConfig +from vllm.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + GateLinear, + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, + is_conv_state_dim_first, +) +from vllm.model_executor.layers.mhc import ( + MHCFusedPostPreOp, + MHCPostOp, + MHCPreOp, + hc_contract, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + scaled_dequantize, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.glm4_1v import ( + Glm4vDummyInputsBuilder, + Glm4vForConditionalGeneration, +) +from vllm.model_executor.models.interfaces import ( + HasInnerState, + IsHybrid, + MixtureOfExperts, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + init_vllm_registered_model, + is_pp_missing_parameter, + make_layers, + maybe_prefix, + sequence_parallel_chunk, +) +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_reduce_scatter, + sp_shard, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig + +from .attention import Glm5NextMLAAttention +from .kda import Glm5NextLinearAttention +from .multimodal import ( + Glm5NextMultiModalProcessor, + Glm5NextProcessingInfo, + Glm5NextVisionTransformer, +) + +logger = init_logger(__name__) + + +def _get_moe_router_dtype(config: Glm5NextConfig) -> torch.dtype | None: + if getattr(config, "moe_router_dtype", None) == "float32": + return torch.float32 + return None + + +class Glm5NextMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + is_sequence_parallel=False, + prefix: str = "", + swiglu_limit: float | None = None, + ) -> None: + super().__init__() + + # If is_sequence_parallel, the input and output tensors are sharded + # across the ranks within the tp_group. In this case the weights are + # replicated and no collective ops are needed. + # Otherwise we use standard TP with an allreduce at the end. + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + + self.swiglu_limit = swiglu_limit + if self.swiglu_limit is not None: + self.act_fn = SiluAndMulWithClamp(swiglu_limit=self.swiglu_limit) + else: + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Glm5NextMoE(nn.Module): + def __init__( + self, + config: Glm5NextConfig, + parallel_config: ParallelConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + apply_routed_scale_to_output: bool = False, + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tensor_model_parallel_rank() + + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + + self.ep_group = get_ep_group().device_group + self.ep_rank = get_ep_group().rank_in_group + self.ep_size = self.ep_group.size() + self.n_routed_experts: int = config.n_routed_experts + self.n_shared_experts: int = config.n_shared_experts + + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe + + if config.hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only silu is supported for now." + ) + + self.router_dtype = _get_moe_router_dtype(config) + self.gate = GateLinear( + config.hidden_size, + config.n_routed_experts, + out_dtype=self.router_dtype, + prefix=f"{prefix}.gate", + ) + if getattr(config, "topk_method", None) == "noaux_tc": + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32) + ) + else: + self.gate.e_score_correction_bias = None + + # Load balancing settings. + eplb_config = parallel_config.eplb_config + self.enable_eplb = parallel_config.enable_eplb + + self.n_redundant_experts = eplb_config.num_redundant_experts + self.n_logical_experts = self.n_routed_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts + self.physical_expert_end = ( + self.physical_expert_start + self.n_local_physical_experts + ) + + swiglu_limit = getattr(config, "swiglu_limit", None) + if config.n_shared_experts is None: + self.shared_experts = None + else: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + + self.shared_experts = Glm5NextMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + is_sequence_parallel=self.is_sequence_parallel, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + swiglu_limit=swiglu_limit, + ) + + self.experts = FusedMoE( + shared_experts=self.shared_experts, + gate=self.gate, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_token, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=getattr(config, "norm_topk_prob", True), + quant_config=quant_config, + use_grouped_topk=True, + num_expert_group=getattr(config, "n_group", 1), + topk_group=getattr(config, "topk_group", 1), + prefix=f"{prefix}.experts", + scoring_func=getattr(config, "scoring_func", "softmax"), + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_output=apply_routed_scale_to_output, + e_score_correction_bias=self.gate.e_score_correction_bias, + enable_eplb=self.enable_eplb, + num_redundant_experts=self.n_redundant_experts, + is_sequence_parallel=self.is_sequence_parallel, + n_shared_experts=None, + router_logits_dtype=self.gate.out_dtype, + swiglu_limit=swiglu_limit, + ) + + def forward( + self, + hidden_states: torch.Tensor, + already_sequence_parallel: bool = False, + ) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + + # Chunk the hidden states so they aren't replicated across TP ranks. + # This avoids duplicate computation in self.experts. + if self.is_sequence_parallel and not already_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + + # The router is always external (self.gate); main's MoERunner expects + # pre-computed router_logits, so compute them here unconditionally. + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + + if self.is_sequence_parallel and not already_sequence_parallel: + final_hidden_states = tensor_model_parallel_all_gather( + final_hidden_states, 0 + ) + final_hidden_states = final_hidden_states[:num_tokens] + + return final_hidden_states.view(num_tokens, hidden_dim) + + +class Glm5NextDecoderLayer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: Glm5NextConfig, + layer_idx: int, + prefix: str = "", + topk_indices_buffer: torch.Tensor | None = None, + is_mtp_layer: bool = False, + **kwargs, + ) -> None: + super().__init__() + + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + self.is_moe = config.is_moe + self.num_hidden_layers = config.num_hidden_layers + self.rms_norm_eps = config.rms_norm_eps + self.num_experts = config.n_routed_experts + self.is_mtp_layer = is_mtp_layer + self.mhc = config.mhc + self.layer_kind = "kda" if config.is_kda_layer(layer_idx) else "mla" + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe + + if config.is_kda_layer(layer_idx): + self.self_attn = Glm5NextLinearAttention( + config=config, + vllm_config=vllm_config, + prefix=f"{prefix}.self_attn", + ) + else: + # MLA layers require the latent head dims, which are guaranteed set + # on MLA configs; narrow away the `int | None`. + assert config.v_head_dim is not None + assert config.kv_lora_rank is not None + self.self_attn = Glm5NextMLAAttention( + vllm_config=vllm_config, + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + max_position_embeddings=config.max_position_embeddings, + cache_config=cache_config, + quant_config=None, # MLA projections are BF16 in checkpoint + prefix=f"{prefix}.self_attn", + topk_indices_buffer=topk_indices_buffer, + skip_rope=getattr(config, "mla_nope", False), + ) + + # MTP layers sit past the base model's hidden layers (layer_idx >= + # num_hidden_layers), so they're outside mlp_layer_types; default them + # to the last base layer's MLP type (sparse/MoE for these checkpoints). + mlp_layer_types = config.mlp_layer_types + mlp_type = ( + mlp_layer_types[layer_idx] + if layer_idx < len(mlp_layer_types) + else (mlp_layer_types[-1] if mlp_layer_types else "sparse") + ) + if self.is_moe and self.num_experts is not None and mlp_type == "sparse": + self.mlp = Glm5NextMoE( + config=config, + parallel_config=parallel_config, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = Glm5NextMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + swiglu_limit=config.swiglu_limit, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + # Cached for the hot forward path (isinstance per layer per step). + self._mlp_is_moe = isinstance(self.mlp, Glm5NextMoE) + # In SP, the attention output projection leaves a partial sum; the + # decoder-layer reduce_scatter after attention completes it (DSv4 pattern). + # MTP layers use the non-mHC path which has no sp_reduce_scatter, so + # their o_proj must still reduce normally. + if self.is_sequence_parallel and not is_mtp_layer: + self.self_attn.o_proj.reduce_results = False + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + if self.mhc and not is_mtp_layer: + # mhc config + self.mhc_num_residual_streams = config.mhc_num_residual_streams + self.mhc_no_norm_weight = config.mhc_no_norm_weight + self.mhc_tau = config.mhc_tau + self.hc_eps = config.hc_eps + self.mhc_sinkhorn_iterations = config.mhc_sinkhorn_iterations + self.mhc_post_mult_value = config.mhc_post_mult_value + + n = config.mhc_num_residual_streams + d_model = n * self.hidden_size + mix_hc = (2 + n) * n + + self.n = n + + # attn hc + self.hc_attn_fn = nn.Parameter( + torch.empty(mix_hc, d_model, dtype=torch.float32) + ) + self.hc_attn_fn_broadcast: torch.Tensor | None = None + self.hc_attn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) + self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + + # ffn hc + self.hc_ffn_fn = nn.Parameter( + torch.empty(mix_hc, d_model, dtype=torch.float32) + ) + self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) + self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + + self.mhc_pre_op = MHCPreOp() + self.mhc_post_op = MHCPostOp() + self.mhc_fused_post_pre_op = MHCFusedPostPreOp() + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None = None, + post: torch.Tensor | None = None, + comb: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: + # 70B or MTP layers: KDA + MoE without HC. + if not self.mhc or self.is_mtp_layer: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + attn_output = self.self_attn( + hidden_states=hidden_states, + positions=positions, + ) + hidden_states, residual = self.post_attention_layernorm( + attn_output, residual=residual + ) + hidden_states = self.mlp(hidden_states) + if self.is_mtp_layer: + # Return the unsummed pair: the MTP caller feeds it straight + # into shared_head's fused_add_rms_norm (one kernel instead of + # a separate residual-add + norm). The sum itself is unchanged + # (fp32-accumulated inside the fused kernel). + return hidden_states, residual, None, None + hidden_states = residual + hidden_states + return hidden_states, residual, None, None + + # mHC start. `post`/`comb` carry the previous layer's deferred + # hc_post inputs (its ffn-pre outputs); when present, fuse that + # hc_post with this layer's attn hc_pre into one kernel (inter-layer + # fusion). Layer 0 has no incoming state -> standalone hc_pre. + x = hidden_states + if post is None: + if self.layer_idx == 0: + if self.hc_attn_fn_broadcast is None: + self.hc_attn_fn_broadcast = ( + self.hc_attn_fn.detach() + .view(-1, self.n, self.hidden_size) + .sum(dim=1) + ) + residual, post, comb, x = torch.ops.vllm.mhc_pre_broadcast_tilelang( + x, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.mhc_post_mult_value, + self.mhc_sinkhorn_iterations, + norm_weight=self.input_layernorm.weight.data, + norm_eps=self.input_layernorm.variance_epsilon, + fn_broadcast=self.hc_attn_fn_broadcast, + ) + else: + residual = x + post, comb, x = self.hc_pre( + x, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + norm_weight=self.input_layernorm.weight.data, + norm_eps=self.input_layernorm.variance_epsilon, + ) + else: + residual, post, comb, x = self.hc_fused_post_pre( + x, + residual, + post, + comb, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + norm_weight=self.input_layernorm.weight.data, + norm_eps=self.input_layernorm.variance_epsilon, + ) + + # Attention needs the full token sequence; mHC above ran on the SP + # shard. Gather for attention, scatter back afterward (DSv4 pattern). + if self.is_sequence_parallel: + x = sp_all_gather(x)[: positions.shape[0]] + + x = self.self_attn( + hidden_states=x, + positions=positions, + ) + + if self.is_sequence_parallel: + x = sp_reduce_scatter(x) + + # Fuse post-attn hc_post + pre-FFN hc_pre (+ RMSNorm) into one kernel. + residual, post, comb, x = self.hc_fused_post_pre( + x, + residual, + post, + comb, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + norm_weight=self.post_attention_layernorm.weight.data, + norm_eps=self.post_attention_layernorm.variance_epsilon, + ) + + # Fully Connected + if self._mlp_is_moe: + x = self.mlp(x, already_sequence_parallel=self.is_sequence_parallel) + else: + x = self.mlp(x) + + # mHC end. The last mHC layer materializes its final hc_post (nothing + # to fuse with) then contracts; every other layer defers its hc_post to + # the next layer's fused pre, returning the state. + if self.layer_idx == self.num_hidden_layers - 1: + x = self.hc_post(x, residual, post, comb) + x = hc_contract(x, self.n) + return x, None, None, None + + return x, residual, post, comb + + def hc_pre( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, + ): + post_mix, res_mix, layer_input = self.mhc_pre_op( + residual=x, + fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + rms_eps=self.rms_norm_eps, + hc_pre_eps=self.hc_eps, + hc_sinkhorn_eps=self.hc_eps, + hc_post_mult_value=self.mhc_post_mult_value, + sinkhorn_repeat=self.mhc_sinkhorn_iterations, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + return post_mix, res_mix, layer_input + + def hc_post( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ): + return self.mhc_post_op(x, residual, post, comb) + + def hc_fused_post_pre( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, + ): + return self.mhc_fused_post_pre_op( + x=x, + residual=residual, + post_layer_mix=post, + comb_res_mix=comb, + fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + rms_eps=self.rms_norm_eps, + hc_pre_eps=self.hc_eps, + hc_sinkhorn_eps=self.hc_eps, + hc_post_mult_value=self.mhc_post_mult_value, + sinkhorn_repeat=self.mhc_sinkhorn_iterations, + n_splits=1, + tile_n=1, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + + +class Glm5NextModel(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + self.config = config + + self.vocab_size = config.vocab_size + self.device = current_platform.device_type + + """ + if config.index_topk is not None: + topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + else: + """ + # `index_topk` is declared on Glm5NextTextConfig with a default of None, + # so hasattr() is True even for full-MLA configs (no kpool indexer). + # Gate on the value being set instead. + self.is_v32 = getattr(config, "index_topk", None) is not None + if self.is_v32: + topk_tokens = config.index_topk + # Reserve room for the incomplete pool tail. + kpool = getattr(config, "index_kpool", 1) or 1 + buffer_width = topk_tokens + (kpool - 1 if kpool > 1 else 0) + # Sparse MLA tiles top-k in 128 columns; padded slots remain masked. + sparse_topk_block_n = 128 + buffer_width = ( + (buffer_width + sparse_topk_block_n - 1) // sparse_topk_block_n + ) * sparse_topk_block_n + topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + buffer_width, + dtype=torch.int32, + device=self.device, + ) + else: + # Full-MLA config (no kpool sparse indexer): no topk buffer. + topk_indices_buffer = None + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + def get_layer(prefix: str): + layer_idx = int(prefix.rsplit(".", 1)[1]) + return Glm5NextDecoderLayer( + vllm_config=vllm_config, + config=config, + layer_idx=layer_idx, + prefix=prefix, + topk_indices_buffer=topk_indices_buffer, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + get_layer, + prefix=f"{prefix}.layers", + ) + # The active slice is fixed after construction; cache it so forward + # doesn't rebuild the slice (a fresh list) every step. + self._active_layers = self.layers[self.start_layer : self.end_layer] + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + self.is_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) + + world_size = get_tensor_model_parallel_world_size() + assert config.num_attention_heads % world_size == 0, ( + "num_attention_heads must be divisible by world_size" + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + hidden = self.config.hidden_size + tensors = { + "hidden_states": torch.zeros( + (batch_size, hidden), dtype=dtype, device=device + ), + } + if not getattr(self.config, "mhc", False): + tensors["residual"] = torch.zeros( + (batch_size, hidden), dtype=dtype, device=device + ) + return IntermediateTensors(tensors) + + n = self.config.mhc_num_residual_streams + tensors.update( + { + "residual": torch.zeros( + (batch_size, n, hidden), dtype=dtype, device=device + ), + "post": torch.zeros( + (batch_size, n, 1), dtype=torch.float32, device=device + ), + "comb": torch.zeros( + (batch_size, n, n), dtype=torch.float32, device=device + ), + } + ) + return IntermediateTensors(tensors) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + post = None + comb = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + post = intermediate_tensors.tensors.get("post") + comb = intermediate_tensors.tensors.get("comb") + + full_num_tokens = positions.shape[0] + if self.is_sequence_parallel and get_pp_group().is_first_rank: + hidden_states = sp_shard(hidden_states) + + for layer in self._active_layers: + hidden_states, residual, post, comb = layer( + positions, hidden_states, residual, post, comb + ) + + if not get_pp_group().is_last_rank: + assert residual is not None + tensors = {"hidden_states": hidden_states, "residual": residual} + if self.config.mhc: + assert post is not None and comb is not None + tensors.update({"post": post, "comb": comb}) + return IntermediateTensors(tensors) + + if self.is_sequence_parallel: + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] + + hidden_states = self.norm(hidden_states) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + # MLA: fuse q_a_proj and kv_a_proj_with_mqa + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + # Indexer: fuse wk and weights_proj + (".wk_weights_proj", ".wk", 0), + (".wk_weights_proj", ".weights_proj", 1), + # KDA: merge q, k, v, b, f_a, g_a projections into one GEMM + (".in_proj_qkvbfg_a", ".q_proj", 0), + (".in_proj_qkvbfg_a", ".k_proj", 1), + (".in_proj_qkvbfg_a", ".v_proj", 2), + (".in_proj_qkvbfg_a", ".b_proj", 3), + (".in_proj_qkvbfg_a", ".f_a_proj", 4), + (".in_proj_qkvbfg_a", ".g_a_proj", 5), + ] + if self.config.is_moe: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts, + ) + else: + expert_params_mapping = [] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + # GLM-5.3-Flash NoPE checkpoints omit the RoPE rows from + # ``kv_a_proj_with_mqa``; pad them with zeros for the model shape. + kv_a_pad_size = 0 + if self.config.mla_nope and self.config.qk_rope_head_dim > 0: + kv_a_pad_size = self.config.qk_rope_head_dim + + _pending_wk_fp8: dict = {} + + for args in weights: + name, loaded_weight = args[:2] + kwargs: dict = args[2] if len(args) > 2 else {} + if "rotary_emb.inv_freq" in name: + continue + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is not None: + continue # skip spec decode layers for main model + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + + # Handle FP8 indexer WK: dequantize to BF16 for fusion with + # weights_proj into wk_weights_proj. + if _try_load_fp8_indexer_wk( + name, + loaded_weight, + _pending_wk_fp8, + params_dict, + loaded_params, + ): + continue + + # FP8 checkpoint: dequantize BF16-kept MLA projections + # (q_a_proj / kv_a_proj_with_mqa / o_proj) to BF16. + if _try_load_fp8_attn_proj( + name, + loaded_weight, + _pending_wk_fp8, + params_dict, + loaded_params, + kv_a_pad_size, + ): + continue + + # Pad kv_a_proj_with_mqa for NoPE models + if kv_a_pad_size > 0 and ".kv_a_proj_with_mqa." in name: + pad = torch.zeros( + kv_a_pad_size, + *loaded_weight.shape[1:], + dtype=loaded_weight.dtype, + device=loaded_weight.device, + ) + loaded_weight = torch.cat([loaded_weight, pad], dim=0) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name_mapped = name.replace(weight_name, param_name) + # QKV fusion: skip if fused module doesn't exist in model + if param_name == ".fused_qkv_a_proj" and name_mapped not in params_dict: + continue + name = name_mapped + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for idx, ( + param_name, + weight_name, + expert_id, + expert_shard_id, + ) in enumerate(expert_params_mapping): + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + expert_id=expert_id, + shard_id=expert_shard_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if ( + name.endswith(".bias") + and name not in params_dict + and not self.config.is_linear_attn + ): # noqa: E501 + continue + # Remapping the name of FP8 kv-scale. + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight, **kwargs) + loaded_params.add(name) + return loaded_params + + +class Glm5NextForCausalLM( + nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid +): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.model_config = vllm_config.model_config + self.vllm_config = vllm_config + self.config = self.model_config.hf_config + quant_config = vllm_config.quant_config + self.quant_config = quant_config + self.model = Glm5NextModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor( + self.config.vocab_size, scale=logit_scale + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs + ) + return hidden_states + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + q_dtype, _, _, recurrent_dtype = MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + ) + return q_dtype, recurrent_dtype + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + q_shape, k_shape, v_shape, recurrent_shape = ( + MambaStateShapeCalculator.kda_state_shape( + tp_size, + hf_config.linear_num_heads, + hf_config.linear_head_dim, + conv_kernel_size=hf_config.linear_conv_kernel_dim, + num_spec=num_spec, + ) + ) + if is_conv_state_dim_first(): + merged_conv_shape = ( + q_shape[0] + k_shape[0] + v_shape[0], + q_shape[1], + ) + else: + merged_conv_shape = ( + q_shape[0], + q_shape[1] + k_shape[1] + v_shape[1], + ) + return merged_conv_shape, recurrent_shape + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + q_copy, _, _, recurrent_copy = ( + MambaStateCopyFuncCalculator.kda_state_copy_func() + ) + return q_copy, recurrent_copy + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + ) + return loader.load_weights(weights) + + +@MULTIMODAL_REGISTRY.register_processor( + Glm5NextMultiModalProcessor, + info=Glm5NextProcessingInfo, + dummy_inputs=Glm4vDummyInputsBuilder, +) +class Glm5NextForConditionalGeneration( + Glm4vForConditionalGeneration, HasInnerState, IsHybrid +): + # The text model (KDA + dense-MLA + MoE) is a hybrid mamba model. The + # multimodal wrapper must declare the same interfaces so vLLM treats it as + # hybrid (auto-aligns mamba/attention block sizes, sizes the mamba state + # cache); the mamba-state classmethods delegate to the text model. + has_inner_state: ClassVar[Literal[True]] = True + is_hybrid: ClassVar[Literal[True]] = True + + # NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration + # (``model.visual.`` -> ``visual.``, ``model.language_model.`` -> + # ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``), + # matching the GLM-OCR / GLM-4V serialization convention. If the real + # checkpoint's safetensors keys differ (e.g. ``language_model.model.`` with + # no outer ``model.``), override ``hf_to_vllm_mapper`` accordingly. + + @classmethod + def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig): + from .model import Glm5NextForCausalLM + + return Glm5NextForCausalLM.get_mamba_state_dtype_from_config(vllm_config) + + @classmethod + def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig): + from .model import Glm5NextForCausalLM + + return Glm5NextForCausalLM.get_mamba_state_shape_from_config(vllm_config) + + @classmethod + def get_mamba_state_copy_func(cls): + from .model import Glm5NextForCausalLM + + return Glm5NextForCausalLM.get_mamba_state_copy_func() + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super(Glm4vForConditionalGeneration, self).__init__() + config = vllm_config.model_config.hf_config + multimodal_config = vllm_config.model_config.multimodal_config + assert multimodal_config is not None + + self.config = config + self.model_config = vllm_config.model_config + self.multimodal_config = multimodal_config + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.visual = Glm5NextVisionTransformer( + config.text_config, + config.vision_config, + # Read eps from the VISION sub-config, not the top-level + # `config.rms_norm_eps`: Glm5NextConfig.__getattribute__ mirrors + # the latter onto text_config (1e-5), silently ignoring the + # vision tower's own (1e-6) rms_norm_eps. + norm_eps=config.vision_config.rms_norm_eps, + # Vision tower ships BF16 weights in this fp8 checkpoint (no + # weight_scale_inv for visual.*), so it must NOT inherit the + # global fp8 quant_config -- doing so incorrectly quantizes + # the tower + # and yields NaN image features. Mirrors the MLA/KDA proj + # pattern (quant_config=None for BF16 submodules). + quant_config=None, + prefix=maybe_prefix(prefix, "visual"), + ) + + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["Glm5NextForCausalLM"], + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + + def get_encoder_cudagraph_config(self): + # The forked vision tower (multimodal.py) has no abs-pos embeddings, so its + # prepare_encoder_metadata does not produce "pos_embeds". Drop it from the + # buffer_keys inherited from Glm4vForConditionalGeneration so encoder + # CUDA-graph capture/replay does not expect a buffer that is never filled. + config = super().get_encoder_cudagraph_config() + config.buffer_keys = [k for k in config.buffer_keys if k != "pos_embeds"] + return config + + +def get_spec_layer_idx_from_weight_name( + config: Glm5NextConfig, weight_name: str +) -> int | None: + if hasattr(config, "num_nextn_predict_layers") and ( + config.num_nextn_predict_layers > 0 + ): + layer_idx = config.num_hidden_layers + for i in range(config.num_nextn_predict_layers): + if weight_name.startswith( + f"model.layers.{layer_idx + i}." + ) or weight_name.startswith(f"layers.{layer_idx + i}."): + return layer_idx + i + return None + + +def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params): + if "indexer.wk." not in name or "wk_weights" in name: + return False + is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn + is_scale = "weight_scale_inv" in name + if not is_weight and not is_scale: + return False + layer_prefix = name.rsplit(".wk.", 1)[0] + entry = buf.setdefault(layer_prefix, {}) + entry["weight" if is_weight else "scale"] = tensor + if "weight" not in entry or "scale" not in entry: + return True + + weight_fp8, scale_inv = entry["weight"], entry["scale"] + del buf[layer_prefix] + block_size = weight_fp8.shape[1] // scale_inv.shape[1] + weight_bf16 = scaled_dequantize( + weight_fp8, + scale_inv, + group_shape=GroupShape(block_size, block_size), + out_dtype=torch.bfloat16, + ) + + fused_name = f"{layer_prefix}.wk_weights_proj.weight" + param = params_dict[fused_name] + param.weight_loader(param, weight_bf16, 0) + loaded_params.add(fused_name) + return True + + +def _dequant_fp8_block( + weight_fp8: torch.Tensor, + scale_inv: torch.Tensor, + block_size: int = 128, +) -> torch.Tensor: + """Dequantize a block-FP8 (e4m3) weight with per-block scale to BF16. + + Unlike ``scaled_dequantize`` this tolerates a non-divisible (partial last + block) shape by zero-padding to a multiple of ``block_size`` before the + scale broadcast and trimming back afterwards (e.g. kv_a_proj_with_mqa is + 576 rows = 4*128 + 64). + """ + out_dim, in_dim = weight_fp8.shape + pad_out = (-out_dim) % block_size + pad_in = (-in_dim) % block_size + w = weight_fp8 + if pad_out or pad_in: + w = torch.nn.functional.pad(w, (0, pad_in, 0, pad_out)) + # scale_inv is (ceil(out/block), ceil(in/block)); broadcast to (out, in). + s = scale_inv.to(torch.float32) + s_full = s.repeat_interleave(block_size, dim=0).repeat_interleave(block_size, dim=1) + out = (w.to(torch.float32) * s_full).to(torch.bfloat16) + return out[:out_dim, :in_dim].contiguous() + + +# FP8 checkpoint projections that the MODEL keeps in BF16, so the block-FP8 +# (weight + weight_scale_inv) must be dequantized to BF16 on load. +# Maps checkpoint proj-suffix -> (buffer key, model target base, fused shard id +# or None for a direct projection, whether NoPE rope-padding applies). +_FP8_ATTN_PROJS = { + ".q_a_proj.": ("q_a", "fused_qkv_a_proj", 0, False), + ".kv_a_proj_with_mqa.": ("kv_a", "fused_qkv_a_proj", 1, True), + ".q_b_proj.": ("q_b", "q_b_proj", None, False), + ".o_proj.": ("o_proj", "o_proj", None, False), +} + + +def _try_load_fp8_attn_proj( + name, + tensor, + buf, + params_dict, + loaded_params, + kv_a_pad_size: int, +) -> bool: + """Dequantize FP8 q_a_proj / kv_a_proj_with_mqa / o_proj to BF16 on load. + + The FP8 checkpoint stores these as block-FP8 (weight + weight_scale_inv), + but the model holds them in BF16 (``fused_qkv_a_proj`` is always BF16 via + DeepSeekV2FusedQkvAProjLinear; ``o_proj`` is excluded by + modules_to_not_convert). When the model target is BF16 (no + ``weight_scale_inv`` param) we dequantize; otherwise we return False so the + normal stacked/direct path loads the FP8 tensor as-is. + """ + matched = None + for suffix, info in _FP8_ATTN_PROJS.items(): + if suffix in name: + matched = (suffix, info) + break + if matched is None: + return False + suffix, (key, target_base, shard_id, is_kva) = matched + is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn + is_scale = "weight_scale_inv" in name + if not is_weight and not is_scale: + return False + + layer_prefix = name.rsplit(suffix, 1)[0] + target_w = f"{layer_prefix}.{target_base}.weight" + target_s = f"{layer_prefix}.{target_base}.weight_scale_inv" + # If the model actually kept this projection in FP8, let the normal path + # handle it (it has a weight_scale_inv param). + if target_s in params_dict: + return False + + entry = buf.setdefault(layer_prefix, {}).setdefault(key, {}) + entry["weight" if is_weight else "scale"] = tensor + if "weight" not in entry or "scale" not in entry: + return True + + weight_fp8, scale_inv = entry["weight"], entry["scale"] + buf[layer_prefix].pop(key, None) + block_size = weight_fp8.shape[1] // scale_inv.shape[1] + weight_bf16 = _dequant_fp8_block(weight_fp8, scale_inv, block_size) + # NoPE: pad kv_a rope portion (kv_lora_rank -> kv_lora_rank + qk_rope_head_dim). + if is_kva and kv_a_pad_size > 0: + pad = torch.zeros( + kv_a_pad_size, + weight_bf16.shape[1], + dtype=weight_bf16.dtype, + device=weight_bf16.device, + ) + weight_bf16 = torch.cat([weight_bf16, pad], dim=0) + + param = params_dict[target_w] + if shard_id is None: + param.weight_loader(param, weight_bf16) + else: + param.weight_loader(param, weight_bf16, shard_id) + loaded_params.add(target_w) + return True diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/nvidia/mtp.py new file mode 100644 index 0000000000..20c3ac8c3d --- /dev/null +++ b/vllm/models/glm5next/nvidia/mtp.py @@ -0,0 +1,424 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import typing +from collections.abc import Callable, Iterable + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.deepseek_mtp import SharedHead +from vllm.model_executor.models.deepseek_v2 import DeepseekV2MixtureOfExperts +from vllm.model_executor.models.utils import maybe_prefix +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors + +from .model import ( + Glm5NextDecoderLayer, + Glm5NextMoE, + _try_load_fp8_attn_proj, + _try_load_fp8_indexer_wk, + get_spec_layer_idx_from_weight_name, +) +from .ops.fused_eh_norm import fused_eh_norm + + +class Glm5NextMultiTokenPredictorLayer(nn.Module): + def __init__(self, vllm_config: VllmConfig, prefix: str) -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + quant_config = vllm_config.quant_config + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + + # Reserve room for the incomplete pool tail and align the sparse MLA + # buffer width to BLOCK_N=128. + topk_tokens = config.index_topk + kpool = getattr(config, "index_kpool", 1) or 1 + buffer_width = topk_tokens + (kpool - 1 if kpool > 1 else 0) + sparse_topk_block_n = 128 + buffer_width = ( + (buffer_width + sparse_topk_block_n - 1) // sparse_topk_block_n + ) * sparse_topk_block_n + topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + buffer_width, + dtype=torch.int32, + device=current_platform.device_type, + ) + self.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=quant_config + ) + # MTP layers sit past the base model's hidden layers; parse the index + # from the prefix (e.g. "...layers.32") so the decoder builds an MLA + # (DSA) layer rather than KDA for the MTP path. + layer_idx = int(prefix.rsplit(".", 1)[-1]) + self.mtp_block = Glm5NextDecoderLayer( + vllm_config=vllm_config, + config=config, + layer_idx=layer_idx, + prefix=prefix, + topk_indices_buffer=topk_indices_buffer, + is_mtp_layer=True, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> torch.Tensor: + assert inputs_embeds is not None + # Fused: zero pos-0 embeds + enorm(embeds) + hnorm(prev) + cat -> [N, 2H]. + eh_input = fused_eh_norm( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + hidden_states = self.eh_proj(eh_input) + # Fuse the residual add and final RMSNorm. Glm5NextMoE already performs + # its all-reduce, so no collective is needed here. The post-norm result + # feeds both draft logits and the next recycled hidden state. + hidden_states, residual, _, _ = self.mtp_block( + positions=positions, hidden_states=hidden_states, residual=None + ) + hidden_states, _ = self.shared_head.norm(hidden_states, residual=residual) + return hidden_states, hidden_states + + +class Glm5NextMultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + self.layers = torch.nn.ModuleDict( + { + str(idx): Glm5NextMultiTokenPredictorLayer( + vllm_config, f"{prefix}.layers.{idx}" + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + # Plain list for the per-propose lookup: ModuleDict[str(...)] builds a + # string and hashes it on every draft step. + self._mtp_layers = list(self.layers.values()) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def set_skip_topk(self, skip: bool): + # index_share_for_mtp_iteration: step 0 computes top-k, steps 1+ reuse. + for layer in self.layers.values(): + self_attn = getattr(layer.mtp_block, "self_attn", None) + if self_attn is not None and hasattr(self_attn, "skip_topk"): + self_attn.skip_topk = skip + + def compact_topk_indices(self, slot_ids: torch.Tensor): + """Gather the top-k index rows at ``slot_ids`` to the front of the buffer.""" + num_slots = slot_ids.numel() + for layer in self.layers.values(): + self_attn = getattr(layer.mtp_block, "self_attn", None) + if self_attn is not None and hasattr(self_attn, "topk_indices_buffer"): + topk_indices_buffer = self_attn.topk_indices_buffer + topk_indices_buffer[:num_slots] = topk_indices_buffer[slot_ids] + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self._mtp_layers[current_step_idx]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self._mtp_layers[current_step_idx] + # hidden_states is already post-final-norm (produced in the layer + # forward and recycled as-is); apply the LM head only, without a + # second RMSNorm. + return self.logits_processor(mtp_layer.shared_head.head, hidden_states) + + def get_top_tokens( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self._mtp_layers[current_step_idx] + # Vocab-parallel argmax for the greedy draft: per-rank head projection + # + local argmax + a [batch, 2*tp] (value, index) reduce, instead of + # materializing and all-gathering full [N, vocab] logits per draft + # step. Tie-breaking matches the full argmax (shards are contiguous + # and rank-ordered, so the lowest-rank winner is the lowest global + # index), so greedy draft tokens are unchanged. + return self.logits_processor.get_top_tokens( + mtp_layer.shared_head.head, hidden_states + ) + + +class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config + self.quant_config = vllm_config.quant_config + self.model = Glm5NextMultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.set_moe_parameters() + + def set_moe_parameters(self): + self.num_moe_layers = self.config.num_nextn_predict_layers + self.num_expert_groups = self.config.n_group + self.moe_layers = [] + self.moe_mlp_layers = [] + example_moe = None + for layer in self.model.layers.values(): + mlp = layer.mtp_block.mlp + if isinstance(mlp, Glm5NextMoE): + example_moe = mlp + self.moe_mlp_layers.append(mlp) + self.moe_layers.append(mlp.experts) + self.extract_moe_parameters(example_moe) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model( + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def get_top_tokens( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + # Greedy-draft path used when use_local_argmax_reduction is enabled: + # vocab-parallel argmax, no full-vocab logits. + return self.model.get_top_tokens(hidden_states, spec_step_idx) + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "eh_proj", + "shared_head", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + name = name.replace( + f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block." + ) + elif shared_weight: + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ("fused_qkv_a_proj", "q_a_proj", 0), + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), + ] + expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts, + ) + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + _pending_wk_fp8: dict = {} + # GLM-5.3-Flash NoPE checkpoints omit the RoPE rows from + # ``kv_a_proj_with_mqa``; the FP8-to-BF16 path pads them for the model. + kv_a_pad_size = 0 + if self.config.mla_nope and self.config.qk_rope_head_dim > 0: + kv_a_pad_size = self.config.qk_rope_head_dim + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + # Multimodal (Glm5NextForConditionalGeneration) checkpoints prefix + # the text-tower weights with "model.language_model."; the MTP head + # is built as a text-only model (model.layers.*), so strip the + # prefix to match. + if name.startswith("model.language_model."): + name = name.replace("model.language_model.", "model.", 1) + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + name = self._rewrite_spec_layer_name(spec_layer, name) + + if _try_load_fp8_indexer_wk( + name, + loaded_weight, + _pending_wk_fp8, + params_dict, + loaded_params, + ): + continue + + # FP8 checkpoint: dequantize the BF16-kept MLA projections + # (q_a_proj / kv_a_proj_with_mqa / o_proj) to BF16, mirroring the + # target model. The model holds fused_qkv_a_proj / o_proj in BF16, + # so the checkpoint's block-FP8 weight + weight_scale_inv for these + # has no param home and would KeyError without this dequant. + if _try_load_fp8_attn_proj( + name, + loaded_weight, + _pending_wk_fp8, + params_dict, + loaded_params, + kv_a_pad_size, + ): + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if ("mlp.experts." in name) and name not in params_dict: + continue + name_mapped = name.replace(weight_name, param_name) + if ( + param_name == "fused_qkv_a_proj" + ) and name_mapped not in params_dict: + continue + else: + name = name_mapped + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment] + if weight_name not in name: + continue + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + param = params_dict[name_mapped] + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + name = name_mapped + break + else: + if is_expert_weight: + continue + if name.endswith(".bias") and name not in params_dict: + continue + name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment] + if name is None: + continue + if ( + spec_layer != self.model.mtp_start_layer_idx + and ".layers" not in name + ): + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_layers.add(spec_layer) + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if layer_idx not in loaded_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint." + ) + return loaded_params diff --git a/vllm/models/glm5next/nvidia/multimodal.py b/vllm/models/glm5next/nvidia/multimodal.py new file mode 100644 index 0000000000..293dffe9fb --- /dev/null +++ b/vllm/models/glm5next/nvidia/multimodal.py @@ -0,0 +1,733 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-5.3-Flash vision tower and multimodal processor.""" + +from collections.abc import Mapping +from functools import partial + +import numpy as np +import torch +import torch.nn as nn +from einops import rearrange + +from vllm.distributed import ( + get_tensor_model_parallel_world_size, + parallel_state, +) +from vllm.distributed import utils as dist_utils +from vllm.model_executor.layers.activation import SiluAndMulWithClamp +from vllm.model_executor.layers.attention import MMEncoderAttention +from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb +from vllm.model_executor.models.glm4_1v import ( + Glm4vMultiModalProcessor, + Glm4vProcessingInfo, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, +) +from vllm.model_executor.models.vision import ( + get_vit_attn_backend, + is_vit_use_data_parallel, +) +from vllm.models.deepseek_v4.common.ops import fused_q_kv_rmsnorm +from vllm.multimodal.parse import ImageSize, MultiModalDataItems +from vllm.v1.attention.backends.registry import AttentionBackendEnum + + +class Glm5NextVisionPatchEmbed(nn.Module): + def __init__( + self, + patch_size: int = 14, + temporal_patch_size: int = 1, + in_channels: int = 3, + hidden_size: int = 1536, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.temporal_patch_size = temporal_patch_size + self.hidden_size = hidden_size + + kernel_size = (temporal_patch_size, patch_size, patch_size) + self.proj = Conv3dLayer( + in_channels, + hidden_size, + kernel_size=kernel_size, + stride=kernel_size, + bias=True, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + L, C = x.shape + x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size) + x = self.proj(x).view(L, self.hidden_size) + return x + + +class Glm5NextVisionMLP(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: int, + swiglu_limit: float, + bias: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + use_data_parallel = is_vit_use_data_parallel() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=in_features, + output_sizes=[hidden_features] * 2, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + disable_tp=use_data_parallel, + ) + self.down_proj = RowParallelLinear( + hidden_features, + in_features, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + disable_tp=use_data_parallel, + ) + # GLM-5.3-Flash clamps the vision SwiGLU gate/up unlike GLM-OCR/GLM-4V. + self.act_fn = SiluAndMulWithClamp(swiglu_limit=swiglu_limit) + + def forward(self, x: torch.Tensor): + x, _ = self.gate_up_proj(x) + x = self.act_fn(x) + x, _ = self.down_proj(x) + return x + + +class Glm5NextVisionAttention(nn.Module): + def __init__( + self, + embed_dim: int, + num_heads: int, + projection_size: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + use_data_parallel = is_vit_use_data_parallel() + self.tp_size = ( + 1 if use_data_parallel else get_tensor_model_parallel_world_size() + ) + self.tp_rank = ( + 0 if use_data_parallel else parallel_state.get_tensor_model_parallel_rank() + ) + self.hidden_size_per_attention_head = dist_utils.divide( + projection_size, num_heads + ) + self.num_attention_heads_per_partition = dist_utils.divide( + num_heads, self.tp_size + ) + + self.head_dim = embed_dim // num_heads + + # q/k norm eps hard-coded 1e-5 — distinct from block/post norm eps. + self.q_norm = RMSNorm(self.head_dim, eps=1e-5) + self.k_norm = RMSNorm(self.head_dim, eps=1e-5) + + self.qkv = QKVParallelLinear( + hidden_size=embed_dim, + head_size=self.hidden_size_per_attention_head, + total_num_heads=num_heads, + total_num_kv_heads=num_heads, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj" if quant_config else f"{prefix}.qkv", + disable_tp=use_data_parallel, + ) + self.proj = RowParallelLinear( + input_size=projection_size, + output_size=embed_dim, + quant_config=quant_config, + prefix=f"{prefix}.proj", + bias=True, + disable_tp=use_data_parallel, + ) + + self.attn = MMEncoderAttention( + num_heads=self.num_attention_heads_per_partition, + head_size=self.hidden_size_per_attention_head, + scale=self.hidden_size_per_attention_head**-0.5, + prefix=f"{prefix}.attn", + ) + self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) + + def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]: + seq_len, bs, _ = qkv.shape + q, k, v = qkv.chunk(3, dim=2) + new_shape = ( + seq_len, + bs, + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ) + q, k, v = (x.view(*new_shape) for x in (q, k, v)) + return q, k, v + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + max_seqlen: torch.Tensor | None = None, + ) -> torch.Tensor: + x, _ = self.qkv(x) + q, k, v = self.split_qkv(x) + + # P1: fused q/k RMSNorm (two distinct weights, one launch; fp32, bit-identical). + q_shape, k_shape = q.shape, k.shape + q_flat = q.reshape(-1, self.head_dim) + k_flat = k.reshape(-1, self.head_dim) + q, k = fused_q_kv_rmsnorm( + q_flat, + k_flat, + self.q_norm.weight, + self.k_norm.weight, + self.q_norm.variance_epsilon, + ) + q = q.view(q_shape) + k = k.view(k_shape) + + q, k, v = (rearrange(t, "s b ... -> b s ...").contiguous() for t in (q, k, v)) + if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None: + qk_concat = torch.cat([q, k], dim=0) + qk_rotated = self.apply_rotary_emb( + qk_concat, + rotary_pos_emb_cos, + rotary_pos_emb_sin, + ) + q, k = torch.chunk(qk_rotated, 2, dim=0) + + context_layer = self.attn( + query=q, + key=k, + value=v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + context_layer = rearrange(context_layer, "b s h d -> s b (h d)").contiguous() + + output, _ = self.proj(context_layer) + return output + + +class Glm5NextVisionBlock(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_hidden_dim: int, + swiglu_limit: float, + norm_layer: partial[nn.Module] | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + self.norm1 = norm_layer(dim) + self.norm2 = norm_layer(dim) + self.attn = Glm5NextVisionAttention( + embed_dim=dim, + num_heads=num_heads, + projection_size=dim, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + self.mlp = Glm5NextVisionMLP( + dim, + mlp_hidden_dim, + swiglu_limit=swiglu_limit, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + max_seqlen: int | None = None, + ) -> torch.Tensor: + x_attn = self.attn( + self.norm1(x), + cu_seqlens=cu_seqlens, + rotary_pos_emb_cos=rotary_pos_emb_cos, + rotary_pos_emb_sin=rotary_pos_emb_sin, + max_seqlen=max_seqlen, + ) + x_fused_norm, residual = self.norm2(x, residual=x_attn) + x = residual + self.mlp(x_fused_norm) + return x + + +class Glm5NextPatchMerger(nn.Module): + def __init__( + self, + d_model: int, + context_dim: int, + swiglu_limit: float, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + use_data_parallel = is_vit_use_data_parallel() + self.hidden_size = d_model + self.proj = ColumnParallelLinear( + self.hidden_size, + self.hidden_size, + bias=bias, + gather_output=True, + quant_config=quant_config, + prefix=f"{prefix}.proj", + disable_tp=use_data_parallel, + ) + self.post_projection_norm = nn.LayerNorm(self.hidden_size) + self.gate_up_proj = MergedColumnParallelLinear( + input_size=self.hidden_size, + output_sizes=[context_dim] * 2, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + disable_tp=use_data_parallel, + ) + self.down_proj = RowParallelLinear( + context_dim, + self.hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + disable_tp=use_data_parallel, + ) + # GLM-5.3-Flash also clamps the merger SwiGLU. + self.act_fn = SiluAndMulWithClamp(swiglu_limit=swiglu_limit) + self.extra_activation_func = nn.GELU() + + def forward(self, x: torch.Tensor): + x, _ = self.proj(x) + x = self.extra_activation_func(self.post_projection_norm(x)) + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Glm5NextVisionTransformer(nn.Module): + # Stacked-weight remap for the GLM-OCR/GLM-4V vision checkpoint layout. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + ".attn.q.": (".attn.qkv.", "q"), + ".attn.k.": (".attn.qkv.", "k"), + ".attn.v.": (".attn.qkv.", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + } + ) + + def __init__( + self, + text_config, # noqa: ANN001 + vision_config, + norm_eps: float = 1e-6, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + use_data_parallel = is_vit_use_data_parallel() + self.tp_size = ( + 1 if use_data_parallel else get_tensor_model_parallel_world_size() + ) + + patch_size = vision_config.patch_size + temporal_patch_size = vision_config.temporal_patch_size + in_channels = vision_config.in_channels + depth = vision_config.depth + self.hidden_size = vision_config.hidden_size + self.num_heads = vision_config.num_heads + + self.patch_size = vision_config.patch_size + self.spatial_merge_size = vision_config.spatial_merge_size + self.out_hidden_size = vision_config.out_hidden_size + + # Fall back to the text clamp limit when vision_config omits it. + swiglu_limit = getattr(vision_config, "swiglu_limit", None) + if swiglu_limit is None: + swiglu_limit = getattr(text_config, "swiglu_limit", None) + assert swiglu_limit is not None, ( + "GLM-5.3-Flash vision requires swiglu_limit (vision_config or text_config)" + ) + + # Single construction pass — no abs-pos embeddings / post-conv norm (OCR delta). + self.patch_embed = Glm5NextVisionPatchEmbed( + patch_size=patch_size, + temporal_patch_size=temporal_patch_size, + in_channels=in_channels, + hidden_size=self.hidden_size, + ) + + norm_layer = partial(RMSNorm, eps=norm_eps) + head_dim = self.hidden_size // self.num_heads + self.rotary_pos_emb = get_rope( + head_size=head_dim, + max_position=8192, + is_neox_style=True, + rope_parameters={"partial_rotary_factor": 0.5}, + ) + self.blocks = nn.ModuleList( + [ + Glm5NextVisionBlock( + dim=self.hidden_size, + num_heads=self.num_heads, + mlp_hidden_dim=vision_config.intermediate_size, + swiglu_limit=swiglu_limit, + norm_layer=norm_layer, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{layer_idx}", + ) + for layer_idx in range(depth) + ] + ) + # GLM-5.3-Flash merger bottleneck width. + self.merger = Glm5NextPatchMerger( + d_model=vision_config.out_hidden_size, + context_dim=vision_config.projection_intermediate_size, + swiglu_limit=swiglu_limit, + quant_config=quant_config, + bias=False, + prefix=f"{prefix}.merger", + ) + + self.downsample = Conv2dLayer( + in_channels=vision_config.hidden_size, + out_channels=vision_config.out_hidden_size, + kernel_size=vision_config.spatial_merge_size, + stride=vision_config.spatial_merge_size, + ) + self.post_layernorm = RMSNorm( + vision_config.hidden_size, eps=vision_config.rms_norm_eps + ) + + self.attn_backend = get_vit_attn_backend( + head_size=head_dim, + dtype=torch.get_default_dtype(), + ) + + @property + def dtype(self) -> torch.dtype: + return self.patch_embed.proj.weight.dtype + + @property + def device(self) -> torch.device: + return self.patch_embed.proj.weight.device + + def rot_pos_emb( + self, grid_thw: list[list[int]] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + pos_ids = [] + for t, h, w in grid_thw: + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + hpos_ids = ( + hpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + .permute(0, 2, 1, 3) + .flatten() + ) + wpos_ids = ( + wpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + .permute(0, 2, 1, 3) + .flatten() + ) + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = max(max(h, w) for _, h, w in grid_thw) + + cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size) + + pos_ids = pos_ids.to(cos.device, non_blocking=True) + cos_combined = cos[pos_ids].flatten(1) + sin_combined = sin[pos_ids].flatten(1) + return cos_combined, sin_combined, pos_ids + + def compute_attn_mask_seqlen( + self, + cu_seqlens: torch.Tensor, + ) -> torch.Tensor | None: + max_seqlen = None + if self.attn_backend in { + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.ROCM_AITER_FA, + AttentionBackendEnum.TRITON_ATTN, + }: + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() + return max_seqlen + + def prepare_encoder_metadata( + self, + grid_thw_list: list[list[int]], + *, + max_batch_size: int | None = None, + max_frames_per_batch: int | None = None, + max_seqlen_override: int | None = None, + device: torch.device | None = None, + ) -> dict[str, torch.Tensor | None]: + """Compute encoder metadata for eager and CUDA graph execution.""" + if device is None: + device = self.device + + metadata: dict[str, torch.Tensor | None] = {} + + rotary_cos, rotary_sin, _ = self.rot_pos_emb(grid_thw_list) + metadata["rotary_pos_emb_cos"] = rotary_cos + metadata["rotary_pos_emb_sin"] = rotary_sin + + grid_thw_np = np.array(grid_thw_list, dtype=np.int32) + patches_per_frame = grid_thw_np[:, 1] * grid_thw_np[:, 2] + cu_seqlens = np.repeat(patches_per_frame, grid_thw_np[:, 0]).cumsum( + dtype=np.int32 + ) + cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens]) + + pad_to = ( + max_frames_per_batch if max_frames_per_batch is not None else max_batch_size + ) + if pad_to is not None: + num_seqs = len(cu_seqlens) - 1 + if num_seqs < pad_to: + cu_seqlens = np.concatenate( + [ + cu_seqlens, + np.full( + pad_to - num_seqs, + cu_seqlens[-1], + dtype=np.int32, + ), + ] + ) + + metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens( + self.attn_backend, cu_seqlens, device + ) + + if max_seqlen_override is not None: + max_seqlen_val = max_seqlen_override + else: + max_seqlen_val = MMEncoderAttention.compute_max_seqlen( + self.attn_backend, cu_seqlens + ) + metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32) + + metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens( + self.attn_backend, + cu_seqlens, + self.hidden_size, + self.tp_size, + device, + ) + + return metadata + + def forward( + self, + x: torch.Tensor, + grid_thw: torch.Tensor | list[list[int]], + *, + encoder_metadata: dict[str, torch.Tensor] | None = None, + ) -> torch.Tensor: + # patchify + x = x.to(device=self.device, dtype=self.dtype) + x = self.patch_embed(x) + + if encoder_metadata is not None: + # Encoder CUDA-graph path (PR #49852): rotary/cu_seqlens/max_seqlen are + # precomputed by prepare_encoder_metadata (which uses rot_pos_emb exactly + # as the eager rebuild does), so reuse them and skip the per-call CPU + # rebuild (the low-GPU-util culprit on multimodal workloads). + rotary_pos_emb_cos = encoder_metadata["rotary_pos_emb_cos"] + rotary_pos_emb_sin = encoder_metadata["rotary_pos_emb_sin"] + cu_seqlens = encoder_metadata["cu_seqlens"] + max_seqlen = encoder_metadata["max_seqlen"] + else: + if isinstance(grid_thw, list): + grid_thw = torch.tensor(grid_thw, dtype=torch.int32) + rotary_pos_emb_cos, rotary_pos_emb_sin, _ = self.rot_pos_emb(grid_thw) + cu_seqlens = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] + ).cumsum(dim=0, dtype=torch.int32) + cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens]) + cu_seqlens = cu_seqlens.to(self.device, non_blocking=True) + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens) + + # transformers + x = x.unsqueeze(1) + for blk in self.blocks: + x = blk( + x, + cu_seqlens=cu_seqlens, + rotary_pos_emb_cos=rotary_pos_emb_cos, + rotary_pos_emb_sin=rotary_pos_emb_sin, + max_seqlen=max_seqlen, + ) + + # adapter + x = self.post_layernorm(x) + x = x.view(-1, self.spatial_merge_size, self.spatial_merge_size, x.shape[-1]) + x = x.permute(0, 3, 1, 2) + x = self.downsample(x).view(-1, self.out_hidden_size) + x = self.merger(x) + return x + + def load_weights(self, weights) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +class Glm5NextProcessingInfo(Glm4vProcessingInfo): + """Wires up the vLLM-native processor for the multimodal checkpoint. + + The checkpoint's ``processor_config.json`` declares a custom ``processor_class`` + and stores its image/video processor configs inline (no standalone + ``preprocessor_config.json``), so ``AutoProcessor`` cannot resolve the + config. We bypass it and build our own ``Glm5NextProcessor`` + (``vllm/transformers_utils/processors/glm5next.py``), a port of the + training-side pipeline that no longer imports transformers' GLM processor + classes. The port applies ``patch_expand_factor`` (checkpoint ships 1) + inside ``smart_resize``'s spatial factor. + """ + + def get_hf_processor(self, **kwargs: object): + proc = getattr(self, "_glm5_hf_processor", None) + if proc is None: + from vllm.transformers_utils.processors.glm5next import Glm5NextProcessor + + proc = Glm5NextProcessor.from_pretrained(self.ctx.model_config.model) + self._glm5_hf_processor = proc + return proc + + def _processor_pixel_budget(self, proc) -> tuple[int, int]: + from vllm.transformers_utils.processors.glm5next import _pixel_budget + + return _pixel_budget( + proc.min_image_tokens, + proc.max_image_tokens, + proc.patch_size, + proc.merge_size, + proc.temporal_patch_size, + ) + + def _get_image_max_pixels(self) -> int: + mm_kwargs = self.ctx.get_merged_mm_kwargs({}) + if (override := mm_kwargs.get("max_pixels")) is not None: + return int(override) + return self._processor_pixel_budget(self.get_hf_processor().image_processor)[1] + + def _get_video_max_pixels(self) -> int: + mm_kwargs = self.ctx.get_merged_mm_kwargs({}) + if (override := mm_kwargs.get("max_pixels")) is not None: + return int(override) + return self._processor_pixel_budget(self.get_hf_processor().video_processor)[1] + + def _get_vision_info( + self, + *, + image_width: int, + image_height: int, + num_frames: int = 16, + do_resize: bool = True, + max_image_pixels: int = 28 * 28 * 2 * 30000, + ) -> tuple[ImageSize, int]: + """GLM-5.3-Flash canvas geometry for token budgeting and dummy inputs. + + The inherited Glm4v path resolves the pixel budget from + ``size.longest_edge`` and resizes with GLM-4V's ``smart_resize``. This + checkpoint's ``processor_config.json`` ships the token-budget style + (``min_image_tokens`` / ``max_image_tokens``) with no ``size`` key, and + the alignment factor carries ``patch_expand_factor`` — resolve both + from the vLLM-native processor so profiling matches runtime geometry. + """ + from vllm.transformers_utils.processors.glm5next import smart_resize + + vision_config = self.get_hf_config().vision_config + patch_size = vision_config.patch_size + merge_size = vision_config.spatial_merge_size + temporal_patch_size = vision_config.temporal_patch_size + + image_processor = self.get_hf_processor().image_processor + factor = patch_size * merge_size * image_processor.patch_expand_factor + # Keep the profiling search viable when the caller's budget is below + # one aligned canvas of the requested duration. + max_image_pixels = max(max_image_pixels, temporal_patch_size * factor * factor) + + if do_resize: + t = num_frames if num_frames > temporal_patch_size else temporal_patch_size + resized_height, resized_width = smart_resize( + t=t, + h=image_height, + w=image_width, + t_factor=temporal_patch_size, + h_factor=factor, + w_factor=factor, + min_pixels=1, + max_pixels=max_image_pixels, + ) + preprocessed_size = ImageSize(width=resized_width, height=resized_height) + else: + preprocessed_size = ImageSize(width=image_width, height=image_height) + + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) + grid_t = max(padded_num_frames // temporal_patch_size, 1) + grid_h = preprocessed_size.height // patch_size + grid_w = preprocessed_size.width // patch_size + + num_patches = grid_t * grid_h * grid_w + num_vision_tokens = num_patches // (merge_size**2) + + return preprocessed_size, num_vision_tokens + + +class Glm5NextMultiModalProcessor(Glm4vMultiModalProcessor): + """The vLLM-native ``Glm5NextProcessor`` extracts image/video features + only and passes the prompt text through unchanged, so prompt expansion + (image token repeat, video frame/timestamp structure) is owned by vLLM's + prompt-update machinery — the inherited ``_get_prompt_updates`` builds + the replacement content and the placeholder scan validates against + exactly that.""" + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False diff --git a/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py b/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py new file mode 100644 index 0000000000..095db3917b --- /dev/null +++ b/vllm/models/glm5next/nvidia/ops/fused_eh_norm.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr): + x = x.to(tl.float32) + mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE + rrms = tl.rsqrt(mean_sq + eps) + w = w.to(tl.float32) + return (x * rrms) * w + + +@triton.jit +def _fused_eh_norm_kernel( + pos_ptr, + embeds_ptr, + embeds_stride, + prev_ptr, + prev_stride, + enorm_w_ptr, + hnorm_w_ptr, + eps, + out_ptr, + out_stride, + H: tl.constexpr, + BLOCK: tl.constexpr, +): + """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm + and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out`` + ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat.""" + tok = tl.program_id(0) + off = tl.arange(0, BLOCK) + mask = off < H + + pos = tl.load(pos_ptr + tok) + e = tl.load(embeds_ptr + tok * embeds_stride + off, mask=mask, other=0.0) + e = tl.where(pos == 0, 0.0, e.to(tl.float32)) + ew = tl.load(enorm_w_ptr + off, mask=mask) + e_normed = _rms_norm(e, ew, eps, H) + tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask) + + p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0) + hw = tl.load(hnorm_w_ptr + off, mask=mask) + p_normed = _rms_norm(p, hw, eps, H) + tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask) + + +def fused_eh_norm( + positions: torch.Tensor, + inputs_embeds: torch.Tensor, + previous_hidden: torch.Tensor, + enorm_w: torch.Tensor, + hnorm_w: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Returns cat([enorm(masked embeds), hnorm(prev_hidden)]) -> [N, 2H].""" + n, h = inputs_embeds.shape + out = torch.empty(n, 2 * h, dtype=inputs_embeds.dtype, device=inputs_embeds.device) + _fused_eh_norm_kernel[(n,)]( + positions, + inputs_embeds, + inputs_embeds.stride(0), + previous_hidden, + previous_hidden.stride(0), + enorm_w, + hnorm_w, + eps, + out, + out.stride(0), + h, + triton.next_power_of_2(h), + ) + return out diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py new file mode 100644 index 0000000000..49ee083d4e --- /dev/null +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -0,0 +1,1061 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""kpool (key-pooling) Triton kernels for the sparse-attention indexer. + +The cache stores POOLS (1 entry per ``pool_size`` consecutive tokens) rather +than individual tokens. ``compress_ratio == pool_size`` on the kv_cache_spec +makes the metadata builder emit pool-granular slot_mapping / seq_lens / +cu_seq_lens / page_table for free; this file supplies the compress-write +kernel (replacing ``indexer_k_quant_and_cache``) and the pool-level topk +helpers (select pools -> expand to tokens -> append tail). +""" + +from __future__ import annotations + +import torch + +from vllm.models.deepseek_v4.common.ops.fp8_software import ( + fp32_to_fp8_e4m3fn_bits, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +# The GLM-5.3-Flash indexer head dimension is fixed at 128. +INDEX_HEAD_DIM = 128 + + +def _is_exact_sm70_cuda() -> bool: + return current_platform.is_cuda() and current_platform.is_device_capability((7, 0)) + + +# Hadamard-128 rotation + + +@triton.jit +def _hadamard128_stage(x, GROUPS: tl.constexpr, STRIDE: tl.constexpr): + x3 = tl.reshape(x, (GROUPS, 2, STRIDE)) + x3 = tl.trans(x3, 0, 2, 1) + a, b = tl.split(x3) + x3 = tl.join(a + b, a - b) + x3 = tl.trans(x3, 0, 2, 1) + return tl.reshape(x3, (128,)) + + +@triton.jit +def _hadamard128(x): + x = _hadamard128_stage(x, 64, 1) + x = _hadamard128_stage(x, 32, 2) + x = _hadamard128_stage(x, 16, 4) + x = _hadamard128_stage(x, 8, 8) + x = _hadamard128_stage(x, 4, 16) + x = _hadamard128_stage(x, 2, 32) + x = _hadamard128_stage(x, 1, 64) + return x * 0.08838834764831845 # 1/sqrt(128) + + +def _hadamard128_torch(x: torch.Tensor) -> torch.Tensor: + """Reference / fallback Hadamard-128 on the last dim (must be 128).""" + import math + + n = x.shape[-1] + assert n == 128, f"_hadamard128 expects last dim 128, got {n}" + h = torch.tensor([[1.0, 1.0], [1.0, -1.0]], dtype=torch.float32, device=x.device) + while h.shape[0] < n: + h = torch.cat([torch.cat([h, h], dim=1), torch.cat([h, -h], dim=1)], dim=0) + h = h / math.sqrt(n) + return x @ h + + +@triton.jit +def _fwht_stage(x, N: tl.constexpr, GROUPS: tl.constexpr, STRIDE: tl.constexpr): + # One FWHT butterfly stage on a flat tensor of N = GROUPS*2*STRIDE elems; + # same construction as _hadamard128_stage but with a parametric N so it can + # process BLOCK_R rows at once. + x3 = tl.reshape(x, (GROUPS, 2, STRIDE)) + x3 = tl.trans(x3, 0, 2, 1) + a, b = tl.split(x3) + x3 = tl.join(a + b, a - b) + x3 = tl.trans(x3, 0, 2, 1) + return tl.reshape(x3, (N,)) + + +@triton.jit +def _fwht_quant_kernel( + q_ptr, + qout_ptr, + sout_ptr, + n_rows, + BLOCK_R: tl.constexpr, + SOFTWARE_FP8: tl.constexpr, + MATERIALIZE_FP16: tl.constexpr, +): + """Fused Hadamard-128 rotation + per-row absmax FP8 (ue8m0) quant. + + Each row uses fp32 butterflies and scaling, rounds to bf16, then applies + absmax quantization with a power-of-two scale. + """ + pid = tl.program_id(0) + rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) + rmask = rows < n_rows + offs = tl.arange(0, 128) + x = tl.load( + q_ptr + rows[:, None] * 128 + offs[None, :], mask=rmask[:, None], other=0.0 + ).to(tl.float32) + + # Flatten so each row's 128 lanes stay contiguous: every stage's + # (GROUPS, 2, STRIDE) tiling has 2*STRIDE dividing 128, so pairs never + # straddle a row boundary. GROUPS of each stage scales by BLOCK_R. + N: tl.constexpr = BLOCK_R * 128 + x = tl.reshape(x, (N,)) + x = _fwht_stage(x, N, BLOCK_R * 64, 1) + x = _fwht_stage(x, N, BLOCK_R * 32, 2) + x = _fwht_stage(x, N, BLOCK_R * 16, 4) + x = _fwht_stage(x, N, BLOCK_R * 8, 8) + x = _fwht_stage(x, N, BLOCK_R * 4, 16) + x = _fwht_stage(x, N, BLOCK_R * 2, 32) + x = _fwht_stage(x, N, BLOCK_R, 64) + x = x * 0.08838834764831845 # 1/sqrt(128), exact in fp32 + + # Match the model activation dtype's materialization before quantizing. + if MATERIALIZE_FP16: + x = x.to(tl.float16).to(tl.float32) + else: + x = x.to(tl.bfloat16).to(tl.float32) + x = tl.reshape(x, (BLOCK_R, 128)) + + absmax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-4) + scale = tl.exp2(tl.ceil(tl.log2(absmax * (1.0 / 448.0)))) + y = tl.minimum(tl.maximum(x / scale[:, None], -448.0), 448.0) + + if SOFTWARE_FP8: + y = fp32_to_fp8_e4m3fn_bits(y) + tl.store(qout_ptr + rows[:, None] * 128 + offs[None, :], y, mask=rmask[:, None]) + tl.store(sout_ptr + rows, scale, mask=rmask) + + +def fwht128_quant_fp8(q: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Rotate each 128-wide row by the Hadamard-128 transform, then FP8-quant. + + The fused kernel avoids materializing the rotated tensor and uses an exact + fp32 ``1 / sqrt(128)`` scale before block-128 ue8m0 quantization. + + Args: + q: ``[rows, 128]`` FP16 or BF16 — one head vector per row. + + Returns: + (q_fp8 ``[rows, 128]`` float8_e4m3fn, scale ``[rows, 1]`` float32). + """ + assert q.ndim == 2 and q.shape[1] == 128, q.shape + assert q.dtype in (torch.bfloat16, torch.float16) + assert q.is_contiguous() + n_rows = q.shape[0] + software_fp8 = _is_exact_sm70_cuda() + q_fp8 = torch.empty((n_rows, 128), dtype=torch.float8_e4m3fn, device=q.device) + q_scale = torch.empty((n_rows, 1), dtype=torch.float32, device=q.device) + if n_rows == 0: + return q_fp8, q_scale + BLOCK_R = 32 + grid = (triton.cdiv(n_rows, BLOCK_R),) + qout = q_fp8.view(torch.uint8) if software_fp8 else q_fp8 + _fwht_quant_kernel[grid]( + q, + qout, + q_scale, + n_rows, + BLOCK_R=BLOCK_R, + SOFTWARE_FP8=software_fp8, + MATERIALIZE_FP16=q.dtype == torch.float16, + num_warps=2, + ) + return q_fp8, q_scale + + +@triton.jit +def _fwht_kernel(q_ptr, out_ptr, n_rows, BLOCK_R: tl.constexpr): + pid = tl.program_id(0) + rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) + rmask = rows < n_rows + offs = tl.arange(0, 128) + x = tl.load( + q_ptr + rows[:, None] * 128 + offs[None, :], + mask=rmask[:, None], + other=0.0, + ).to(tl.float32) + + n: tl.constexpr = BLOCK_R * 128 + x = tl.reshape(x, (n,)) + x = _fwht_stage(x, n, BLOCK_R * 64, 1) + x = _fwht_stage(x, n, BLOCK_R * 32, 2) + x = _fwht_stage(x, n, BLOCK_R * 16, 4) + x = _fwht_stage(x, n, BLOCK_R * 8, 8) + x = _fwht_stage(x, n, BLOCK_R * 4, 16) + x = _fwht_stage(x, n, BLOCK_R * 2, 32) + x = _fwht_stage(x, n, BLOCK_R, 64) + x = tl.reshape(x * 0.08838834764831845, (BLOCK_R, 128)) + tl.store( + out_ptr + rows[:, None] * 128 + offs[None, :], + x, + mask=rmask[:, None], + ) + + +def fwht128(q: torch.Tensor) -> torch.Tensor: + """Rotate contiguous 128-wide FP16/BF16 rows without quantizing them.""" + assert q.ndim == 2 and q.shape[1] == 128, q.shape + assert q.dtype in (torch.float16, torch.bfloat16) + assert q.is_contiguous() + out = torch.empty_like(q) + if q.shape[0] == 0: + return out + block_r = 32 + _fwht_kernel[(triton.cdiv(q.shape[0], block_r),)]( + q, + out, + q.shape[0], + BLOCK_R=block_r, + num_warps=2, + ) + return out + + +# Map pool ids to physical cache slots. + + +def compute_pooled_write_locs( + page_table_64: torch.Tensor, + pool_ids: torch.Tensor, + pool_size: int, +) -> torch.Tensor: + """Map logical pooled-K ids to physical flat cache slots. + + ``pool_size`` consecutive tokens share one pool slot that lives at the + *first* token page of each page-group. ``page_table_64`` maps token pages + to physical block ids; we gather the block id of each pool's page-group + and add the in-block pool offset. + """ + assert page_table_64.ndim == 1 + pool_ids = pool_ids.to(torch.int64) + block_size = 64 + pool_page_group = torch.div(pool_ids, block_size, rounding_mode="floor") + token_page_row = pool_page_group * pool_size + packed_page = page_table_64.index_select(0, token_page_row.to(torch.int64)) + return packed_page.to(torch.int64) * block_size + torch.remainder( + pool_ids, block_size + ) + + +def build_pooled_page_table( + page_table: torch.Tensor, + pool_size: int, +) -> torch.Tensor: + """Build a pool-granular page table by taking every ``pool_size``-th + token-page column (one pool maps to ``pool_size`` token pages). + + Uses gather (not strided slicing) so the result is always a fresh + row-major tensor — some downstream kernels require stride(-1) == 1. + """ + block_size = page_table.shape[-1] + assert block_size % pool_size == 0, ( + f"pool_size ({pool_size}) must divide page columns ({block_size})" + ) + idx = torch.arange(0, block_size, pool_size, device=page_table.device) + return page_table[..., idx].contiguous() + + +# Fused pool compression and cache write. + + +@triton.jit +def _kpool_softmax_rotate_write_cache_kernel( + buf_fp8_ptr, + buf_fp32_ptr, + slot_k_ptr, + slot_score_ptr, + ape_ptr, + loc_ptr, + write_mask_ptr, + compressed_k_ptr, + compressed_scale_ptr, + slot_k_stride_0, + slot_k_stride_1, + slot_score_stride_0, + slot_score_stride_1, + ape_stride_0, + PAGE_SIZE: tl.constexpr, + BUF_NUMEL_PER_PAGE: tl.constexpr, + POOL_SIZE: tl.constexpr, + HEAD_DIM: tl.constexpr, + S_OFFSET_NBYTES_IN_PAGE: tl.constexpr, + ROUND_SCALE: tl.constexpr, + HAS_WRITE_MASK: tl.constexpr, + RETURN_COMPRESSED: tl.constexpr, + WRITE_CACHE: tl.constexpr, + SOFTWARE_FP8: tl.constexpr, + MATERIALIZE_FP16: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """One program per pool. softmax(slot_score+ape)-weighted sum of slot_k -> + Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``.""" + row = tl.program_id(0) + do_write = True + if HAS_WRITE_MASK: + do_write = tl.load(write_mask_ptr + row) + + offs = tl.arange(0, BLOCK_D) + mask = (offs < HEAD_DIM) & do_write + + # --- Pass 1: per-dim max over the pool (softmax numerical stability) --- + max_score = tl.full((BLOCK_D,), -float("inf"), tl.float32) + for slot in tl.static_range(0, POOL_SIZE): + score = tl.load( + slot_score_ptr + + row * slot_score_stride_0 + + slot * slot_score_stride_1 + + offs, + mask=mask, + other=0.0, + ).to(tl.float32) + score += tl.load(ape_ptr + slot * ape_stride_0 + offs, mask=mask, other=0.0).to( + tl.float32 + ) + max_score = tl.maximum(max_score, score) + + # --- Pass 2: softmax-weighted sum of K --- + acc = tl.full((BLOCK_D,), 0.0, tl.float32) + denom = tl.full((BLOCK_D,), 0.0, tl.float32) + for slot in tl.static_range(0, POOL_SIZE): + score = tl.load( + slot_score_ptr + + row * slot_score_stride_0 + + slot * slot_score_stride_1 + + offs, + mask=mask, + other=0.0, + ).to(tl.float32) + score += tl.load(ape_ptr + slot * ape_stride_0 + offs, mask=mask, other=0.0).to( + tl.float32 + ) + prob = tl.exp(score - max_score) + denom += prob + k = tl.load( + slot_k_ptr + row * slot_k_stride_0 + slot * slot_k_stride_1 + offs, + mask=mask, + other=0.0, + ).to(tl.float32) + acc += k * prob + + x = acc / denom + x = tl.where(do_write, x, 0.0) + if MATERIALIZE_FP16: + x = x.to(tl.float16).to(tl.float32) + else: + x = x.to(tl.bfloat16).to(tl.float32) + + # Match the unfused pooled-K path's activation precision before quantization. + x = _hadamard128(x) + if MATERIALIZE_FP16: + x = x.to(tl.float16).to(tl.float32) + else: + x = x.to(tl.bfloat16).to(tl.float32) + + # --- per-vector absmax fp8 quant --- + fp8_max = 448.0 + fp8_max_inv = 1.0 / fp8_max + absmax = tl.max(tl.abs(x), axis=0) + absmax = tl.maximum(absmax, 1e-4) + if ROUND_SCALE: + scale = tl.exp2(tl.ceil(tl.log2(absmax * fp8_max_inv))) + else: + scale = absmax * fp8_max_inv + quantized = x / scale + quantized = tl.minimum(tl.maximum(quantized, -fp8_max), fp8_max) + if SOFTWARE_FP8: + quantized = fp32_to_fp8_e4m3fn_bits(quantized) + + if WRITE_CACHE: + loc = tl.load(loc_ptr + row, mask=do_write, other=0) + loc_page_index = loc // PAGE_SIZE + loc_token_offset_in_page = loc % PAGE_SIZE + out_k_offsets = ( + loc_page_index * BUF_NUMEL_PER_PAGE + + loc_token_offset_in_page * HEAD_DIM + + offs + ) + out_s_offset = ( + loc_page_index * BUF_NUMEL_PER_PAGE // 4 + + S_OFFSET_NBYTES_IN_PAGE // 4 + + loc_token_offset_in_page + ) + tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=mask) + tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write) + + if RETURN_COMPRESSED: + tl.store( + compressed_k_ptr + row * HEAD_DIM + offs, + quantized, + mask=offs < HEAD_DIM, + ) + tl.store(compressed_scale_ptr + row, scale) + + +def kpool_compress_and_write_cache( + kv_cache: torch.Tensor, + slot_k: torch.Tensor, + slot_score: torch.Tensor, + ape: torch.Tensor, + loc: torch.Tensor, + pool_size: int, + head_dim: int = INDEX_HEAD_DIM, + write_mask: torch.Tensor | None = None, + round_scale: bool = True, + return_compressed: bool = False, + write_cache: bool = True, +): + """Compress ``pool_size`` tokens into one fp8 K and write at ``loc``. + + Args: + kv_cache: indexer K cache ``[num_blocks, block_size, head_dim+4]`` uint8. + slot_k: ``[n_pools, pool_size, head_dim]`` FP16/BF16 raw per-token K. + slot_score: ``[n_pools, pool_size, head_dim]`` — per-token gate score. + ape: ``[pool_size, head_dim]`` fp32 — per-slot position bias. + loc: ``[n_pools]`` int64 — flat physical slot per pool. + """ + assert slot_k.ndim == 3 + assert slot_score.shape == slot_k.shape + assert ape.shape == slot_k.shape[1:] + assert slot_k.shape[2] == head_dim + assert slot_k.dtype in (torch.bfloat16, torch.float16) + assert slot_score.dtype == slot_k.dtype + assert ape.dtype == torch.float32 + assert kv_cache.dtype == torch.uint8 + assert loc.dtype == torch.int64 + assert write_cache or return_compressed + + page_size = kv_cache.shape[1] + buf = kv_cache + slot_k = slot_k.contiguous() + slot_score = slot_score.contiguous() + ape = ape.contiguous() + loc = loc.contiguous() + if write_mask is None: + write_mask = torch.empty((1,), dtype=torch.bool, device=slot_k.device) + has_write_mask = False + else: + assert write_mask.shape == (slot_k.shape[0],) + write_mask = write_mask.contiguous() + has_write_mask = True + assert not return_compressed + + if slot_k.shape[0] == 0: + if return_compressed: + return ( + torch.empty( + (0, head_dim), + dtype=torch.float8_e4m3fn, + device=slot_k.device, + ), + torch.empty((0,), dtype=torch.float32, device=slot_k.device), + ) + return None + + software_fp8 = _is_exact_sm70_cuda() + buf_fp8 = buf if software_fp8 else buf.view(torch.float8_e4m3fn) + buf_fp32 = buf.view(torch.float32) + # bytes per page (last dim of kv_cache) viewed as uint8 + buf_numel_per_page = buf.stride(0) + s_offset_nbytes_in_page = page_size * head_dim + + if return_compressed: + compressed_k = torch.empty( + (slot_k.shape[0], head_dim), + dtype=torch.float8_e4m3fn, + device=slot_k.device, + ) + compressed_scale = torch.empty( + (slot_k.shape[0],), dtype=torch.float32, device=slot_k.device + ) + else: + compressed_k = buf_fp8 + compressed_scale = buf_fp32 + + compressed_k_ptr = ( + compressed_k.view(torch.uint8) + if software_fp8 and compressed_k.dtype == torch.float8_e4m3fn + else compressed_k + ) + + _kpool_softmax_rotate_write_cache_kernel[(slot_k.shape[0],)]( + buf_fp8, + buf_fp32, + slot_k, + slot_score, + ape, + loc, + write_mask, + compressed_k_ptr, + compressed_scale, + slot_k.stride(0), + slot_k.stride(1), + slot_score.stride(0), + slot_score.stride(1), + ape.stride(0), + PAGE_SIZE=page_size, + BUF_NUMEL_PER_PAGE=buf_numel_per_page, + POOL_SIZE=slot_k.shape[1], + HEAD_DIM=head_dim, + S_OFFSET_NBYTES_IN_PAGE=s_offset_nbytes_in_page, + ROUND_SCALE=round_scale, + HAS_WRITE_MASK=has_write_mask, + RETURN_COMPRESSED=return_compressed, + WRITE_CACHE=write_cache, + SOFTWARE_FP8=software_fp8, + MATERIALIZE_FP16=slot_k.dtype == torch.float16, + BLOCK_D=triton.next_power_of_2(head_dim), + ) + + if return_compressed: + return compressed_k, compressed_scale + return None + + +# Seed each request's incomplete pool into its paged tail during prefill. + + +@triton.jit +def _kpool_tail_seed_kernel( + key_ptr, + score_ptr, + tslot_ptr, + tail_ptr, + n_tokens, + HEAD_DIM: tl.constexpr, + KPOOL: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Copy token ``i``'s raw K + gate into its request's tail block. + + Token ``i`` is among its request's last KPOOL tokens iff the token KPOOL + ahead belongs to a different tail block (or is past the batch / padding, + slot < 0). ``tslot = block * KPOOL + pos % KPOOL``; the destination is + ``tail[block, {0:K, 1:score}, pos % KPOOL, :]``. + """ + i = tl.program_id(0) + t = tl.load(tslot_ptr + i).to(tl.int64) + if t < 0: + return + blk = t // KPOOL # t >= 0 here, so trunc == floor + ahead = tl.load(tslot_ptr + i + KPOOL, mask=i + KPOOL < n_tokens, other=-1).to( + tl.int64 + ) + # Match the torch semantics exactly: a negative ahead slot floors to a + # block id that differs from every real block -> token is in the tail. + # Only divide non-negative slots (Triton int div truncates, torch floors). + if ahead >= 0 and ahead // KPOOL == blk: + return + offs = tl.arange(0, BLOCK_D) + m = offs < HEAD_DIM + base = (blk * 2 * KPOOL + t % KPOOL) * HEAD_DIM + k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m) + s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m) + tl.store(tail_ptr + base + offs, k, mask=m) + tl.store(tail_ptr + base + KPOOL * HEAD_DIM + offs, s, mask=m) + + +def kpool_seed_tail_cache( + tail_kv_cache: torch.Tensor, + key: torch.Tensor, + gate_score: torch.Tensor, + tslot: torch.Tensor, + kpool: int, + head_dim: int = INDEX_HEAD_DIM, +) -> None: + """Seed the paged tail cache from a prefill batch (see the kernel).""" + assert tail_kv_cache.dtype in (torch.bfloat16, torch.float16) + assert key.dtype == tail_kv_cache.dtype + assert gate_score.dtype == key.dtype + n = tslot.shape[0] + if n == 0: + return + _kpool_tail_seed_kernel[(n,)]( + key, + gate_score, + tslot, + tail_kv_cache, + n, + HEAD_DIM=head_dim, + KPOOL=kpool, + BLOCK_D=triton.next_power_of_2(head_dim), + ) + + +# Update each request's tail during decode and write completed pools. + + +@triton.jit +def _kpool_decode_update_batched_kernel( + buf_fp8_ptr, + buf_fp32_ptr, + tail_kv_ptr, + tail_slot_mapping_ptr, # [B, NEXT_N] int32 + key_ptr, # [B, NEXT_N, HEAD_DIM] bf16 + key_stride_b, + key_stride_t, + slot_score_ptr, # [B, NEXT_N, HEAD_DIM] bf16 + ss_stride_b, + ss_stride_t, + ape_ptr, + ape_stride_0, + slot_mapping_ptr, # [B, NEXT_N] int32 + positions_ptr, # [B, NEXT_N] int32 + NEXT_N, # runtime token count per request (no .item() needed) + PAGE_SIZE: tl.constexpr, + BUF_NUMEL_PER_PAGE: tl.constexpr, + POOL_SIZE: tl.constexpr, + TAIL_BLOCK_ELEMS: tl.constexpr, + KPOOL_HEAD: tl.constexpr, + HEAD_DIM: tl.constexpr, + S_OFFSET_NBYTES_IN_PAGE: tl.constexpr, + ROUND_SCALE: tl.constexpr, + SOFTWARE_FP8: tl.constexpr, + MATERIALIZE_FP16: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """One program per request; iterates its NEXT_N verify tokens in order. + + Replaces the caller's per-token sequential launch loop. The intra-request + iteration MUST stay in position order: a pool-completion at token t* reads + the tail-ring slots that tokens t < t* (same request) just stashed in this + same invocation. ``tl.range`` iterates sequentially within the program, so + those stashes are visible to the later completion read. Cross-request + programs are independent (distinct tail blocks). With NEXT_N < POOL_SIZE + (the spec-verify case: NEXT_N ~= num_spec+1, POOL_SIZE=16) at most one + completion can occur per request per call, but the ordered loop is correct + for any NEXT_N. + """ + req = tl.program_id(0) + offs = tl.arange(0, BLOCK_D) + dim_mask = offs < HEAD_DIM + + for t in tl.range(0, NEXT_N): + idx = req * NEXT_N + t + cache_loc = tl.load(slot_mapping_ptr + idx) + pos = tl.load(positions_ptr + idx) + safe_pos = tl.maximum(pos, 0) + pos_valid = (cache_loc >= 0) & (pos >= 0) + + slot = safe_pos % POOL_SIZE + phys_slot = safe_pos % POOL_SIZE + + # Derive the tail block from THIS token's tail_slot (the request's block + # is constant across a pool, but a padded / invalid entry carries a + # negative sentinel -- reading it from token 0 would poison every + # token's base address). Clamp so an invalid entry can never form an + # out-of-bounds base; the accesses below are gated on pos_valid anyway. + tail_slot = tl.load(tail_slot_mapping_ptr + idx) + block = tl.maximum(tail_slot, 0).to(tl.int64) // POOL_SIZE + block_base = block * TAIL_BLOCK_ELEMS + + # The tail-ring stash must run for EVERY real token, so it is gated on + # the token-granular tail slot -- not on `pos_valid`, which keys off the + # POOL-granular `slot_mapping` and is therefore only true on the pool's + # last token. Gating the stash on pos_valid dropped every intra-pool + # token, so a decode-built pool compressed 3 stale ring entries (the + # prefill-seeded prompt tail, frozen forever) plus the current token. + stash_valid = (pos >= 0) & (tail_slot >= 0) + + key = tl.load( + key_ptr + req * key_stride_b + t * key_stride_t + offs, + mask=dim_mask, + other=0.0, + ).to(tl.float32) + score_current = tl.load( + slot_score_ptr + req * ss_stride_b + t * ss_stride_t + offs, + mask=dim_mask, + other=0.0, + ).to(tl.float32) + + if pos_valid & (slot == POOL_SIZE - 1): + pool_logical_start = safe_pos - slot + + max_score = tl.full((BLOCK_D,), -float("inf"), tl.float32) + for pool_slot in tl.static_range(0, POOL_SIZE): + is_current = pool_slot == slot + phys = (pool_logical_start + pool_slot) % POOL_SIZE + score_buf = tl.load( + tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs, + mask=dim_mask, + other=0.0, + ).to(tl.float32) + score = tl.where(is_current, score_current, score_buf) + score += tl.load( + ape_ptr + pool_slot * ape_stride_0 + offs, + mask=dim_mask, + other=0.0, + ).to(tl.float32) + max_score = tl.maximum(max_score, score) + + acc = tl.full((BLOCK_D,), 0.0, tl.float32) + denom = tl.full((BLOCK_D,), 0.0, tl.float32) + for pool_slot in tl.static_range(0, POOL_SIZE): + is_current = pool_slot == slot + phys = (pool_logical_start + pool_slot) % POOL_SIZE + score_buf = tl.load( + tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs, + mask=dim_mask, + other=0.0, + ).to(tl.float32) + score = tl.where(is_current, score_current, score_buf) + score += tl.load( + ape_ptr + pool_slot * ape_stride_0 + offs, + mask=dim_mask, + other=0.0, + ).to(tl.float32) + prob = tl.exp(score - max_score) + denom += prob + k_buf = tl.load( + tail_kv_ptr + block_base + phys * HEAD_DIM + offs, + mask=dim_mask, + other=0.0, + ).to(tl.float32) + k = tl.where(is_current, key, k_buf) + acc += k * prob + + x = acc / denom + if MATERIALIZE_FP16: + x = x.to(tl.float16).to(tl.float32) + else: + x = x.to(tl.bfloat16).to(tl.float32) + x = _hadamard128(x) + if MATERIALIZE_FP16: + x = x.to(tl.float16).to(tl.float32) + else: + x = x.to(tl.bfloat16).to(tl.float32) + + fp8_max = 448.0 + fp8_max_inv = 1.0 / fp8_max + absmax = tl.maximum(tl.max(tl.abs(x), axis=0), 1e-4) + if ROUND_SCALE: + scale = tl.exp2(tl.ceil(tl.log2(absmax * fp8_max_inv))) + else: + scale = absmax * fp8_max_inv + quantized = tl.minimum(tl.maximum(x / scale, -fp8_max), fp8_max) + if SOFTWARE_FP8: + quantized = fp32_to_fp8_e4m3fn_bits(quantized) + + loc = cache_loc.to(tl.int64) + loc_page_index = loc // PAGE_SIZE + loc_token_offset_in_page = loc % PAGE_SIZE + out_k_offsets = ( + loc_page_index * BUF_NUMEL_PER_PAGE + + loc_token_offset_in_page * HEAD_DIM + + offs + ) + out_s_offset = ( + loc_page_index * BUF_NUMEL_PER_PAGE // 4 + + S_OFFSET_NBYTES_IN_PAGE // 4 + + loc_token_offset_in_page + ) + tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=dim_mask) + tl.store(buf_fp32_ptr + out_s_offset, scale) + + # Stash the current token AFTER any completion read so the completion + # uses prior stashes (and the current token's own key/score via + # is_current), then leaves this token for future pools. Order matches + # the per-token kernel: completion read first, stash second. + update_mask = dim_mask & stash_valid + tl.store( + tail_kv_ptr + block_base + phys_slot * HEAD_DIM + offs, + key, + mask=update_mask, + ) + tl.store( + tail_kv_ptr + block_base + KPOOL_HEAD + phys_slot * HEAD_DIM + offs, + score_current, + mask=update_mask, + ) + + +def kpool_decode_update_and_maybe_write_cache_batched( + kv_cache: torch.Tensor, + tail_kv_cache: torch.Tensor, + tail_slot_mapping: torch.Tensor, + key: torch.Tensor, + slot_score: torch.Tensor, + ape: torch.Tensor, + slot_mapping: torch.Tensor, + positions: torch.Tensor, + pool_size: int, + head_dim: int = INDEX_HEAD_DIM, + round_scale: bool = True, +) -> None: + """Batched decode-step kpool update for spec verify (``next_n > 1``). + + One launch replaces the caller's per-token loop. Inputs are grouped per + request: ``[num_requests, next_n, ...]``. Each program handles one + request's ``next_n`` tokens in position order (see the kernel docstring for + why ordering is required for pool-completion correctness). + + Plain decode (``next_n == 1``) is handled here too — the kernel collapses + to a single-iteration loop. + + Args: + kv_cache: indexer K cache ``[num_blocks, block_size, head_dim+4]`` uint8. + tail_kv_cache: paged tail cache ``[num_blocks, 2, pool_size, head_dim]`` + FP16/BF16 (K at half 0, gate score at half 1). + tail_slot_mapping: ``[num_requests, next_n]`` int32. + key / slot_score: ``[num_requests, next_n, head_dim]`` FP16/BF16. + ape: ``[pool_size, head_dim]`` fp32. + slot_mapping / positions: ``[num_requests, next_n]`` int32. + """ + num_requests, next_n = key.shape[0], key.shape[1] + if num_requests == 0 or next_n == 0: + return + assert tail_kv_cache.ndim == 4 + assert tail_kv_cache.shape[1] == 2 + assert tail_kv_cache.shape[2] == pool_size + assert tail_kv_cache.shape[3] == head_dim + assert tail_kv_cache.dtype in (torch.bfloat16, torch.float16) + assert key.ndim == 3 and key.shape[2] == head_dim + assert slot_score.shape == key.shape + assert ape.shape == (pool_size, head_dim) + assert tail_slot_mapping.shape == (num_requests, next_n) + assert slot_mapping.shape == (num_requests, next_n) + assert positions.shape == (num_requests, next_n) + assert key.dtype in (torch.bfloat16, torch.float16) + assert slot_score.dtype == key.dtype + assert tail_kv_cache.dtype == key.dtype + assert ape.dtype == torch.float32 + assert kv_cache.dtype == torch.uint8 + + page_size = kv_cache.shape[1] + buf = kv_cache + software_fp8 = _is_exact_sm70_cuda() + buf_fp8 = buf if software_fp8 else buf.view(torch.float8_e4m3fn) + buf_fp32 = buf.view(torch.float32) + + # The kernel indexes the int tensors as ``req * next_n + t`` (row-major), + # so they must be contiguous. Callers pass either a view of a contiguous + # slice or a freshly scattered tensor, making these no-ops; the calls guard + # against a future caller handing over a strided view. + tail_slot_mapping = tail_slot_mapping.contiguous() + slot_mapping = slot_mapping.contiguous() + positions = positions.contiguous() + + _kpool_decode_update_batched_kernel[(num_requests,)]( + buf_fp8, + buf_fp32, + tail_kv_cache, + tail_slot_mapping, + key, + key.stride(0), + key.stride(1), + slot_score, + slot_score.stride(0), + slot_score.stride(1), + ape, + ape.stride(0), + slot_mapping, + positions, + next_n, + PAGE_SIZE=page_size, + BUF_NUMEL_PER_PAGE=buf.stride(0), + POOL_SIZE=pool_size, + TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0), + KPOOL_HEAD=tail_kv_cache.stride(1), + HEAD_DIM=head_dim, + S_OFFSET_NBYTES_IN_PAGE=page_size * head_dim, + ROUND_SCALE=round_scale, + SOFTWARE_FP8=software_fp8, + MATERIALIZE_FP16=key.dtype == torch.float16, + BLOCK_D=triton.next_power_of_2(head_dim), + ) + + +# Pool-level top-k helpers. + + +def history_group_budget_for_topk(topk: int, pool_size: int) -> int: + """Number of pools to select so that expanding yields ``topk`` tokens.""" + assert topk % pool_size == 0 + return topk // pool_size + + +def expand_pools_to_tokens( + group_ids: torch.Tensor, + group_valid: torch.Tensor, + topk: int, + pool_size: int, + page_table: torch.Tensor | None = None, + topk_offsets: torch.Tensor | None = None, +) -> torch.Tensor: + """Expand selected full-pool ids to a strict-width token topk tensor.""" + assert group_ids.ndim == 2 + assert group_valid.shape == group_ids.shape + assert topk % pool_size == 0 + assert group_ids.shape[1] == history_group_budget_for_topk(topk, pool_size) + assert page_table is None or topk_offsets is None + + device = group_ids.device + offsets = torch.arange(pool_size, device=device, dtype=torch.int64) + token_ids = group_ids.to(torch.int64).unsqueeze(-1) * pool_size + offsets + token_ids = token_ids.reshape(group_ids.shape[0], topk) + valid = ( + group_valid.unsqueeze(-1) + .expand(-1, -1, pool_size) + .reshape(group_ids.shape[0], topk) + ) + + if page_table is not None: + assert page_table.ndim == 2 + safe_ids = token_ids.clamp(min=0, max=page_table.shape[1] - 1) + output = torch.gather(page_table, dim=1, index=safe_ids).to(torch.int32) + elif topk_offsets is not None: + if topk_offsets.ndim == 2: + assert topk_offsets.shape[1] == 1 + topk_offsets = topk_offsets.squeeze(1) + output = (token_ids + topk_offsets.to(torch.int64).unsqueeze(1)).to(torch.int32) + else: + output = token_ids.to(torch.int32) + + return torch.where(valid, output, torch.full_like(output, -1)) + + +def append_tail_to_topk( + topk_result: torch.Tensor, + seq_lens: torch.Tensor, + pool_lens: torch.Tensor, + pool_size: int, + page_table: torch.Tensor | None = None, + topk_offsets: torch.Tensor | None = None, +) -> torch.Tensor: + """Append non-pooled tail tokens after expanded history tokens. + + ``index_kpool_always_select_tail`` keeps the (incomplete) trailing pool so + the most recent tokens are always attended to. + """ + assert topk_result.dtype == torch.int32 + assert seq_lens.ndim == 1 + assert pool_lens.ndim == 1 + + tail_pool = pool_size - 1 + if tail_pool == 0: + return topk_result + + rows, n_cols = topk_result.shape + out_cols = n_cols + tail_pool + out = torch.empty( + (rows, out_cols), dtype=topk_result.dtype, device=topk_result.device + ) + + # tail tokens: [pool_len*pool_size, seq_len) for each row. + pool_len = pool_lens.to(torch.int32) + tail_start = pool_len * pool_size + seq_len = seq_lens.to(torch.int32) + tail_count = seq_len - tail_start # in [0, pool_size) + + cols = torch.arange(out_cols, device=topk_result.device)[None, :] + history_len = n_cols + is_history = cols < history_len + tail_off = cols - history_len + is_tail = (tail_off >= 0) & (tail_off < tail_count[:, None]) + + # safe_hist must be per-row [rows, out_cols] so the gather reads each row's + # OWN history. cols is [1, out_cols]; if used directly, gather (which does + # NOT broadcast the index) would read only row 0 of topk_result, making every + # query inherit row 0's history (empty for the first token) and lose all its + # selected tokens — only the per-row tail would survive. This only manifests + # for multi-row sparse PREFILL (decode has 1 row, so it reads its own row 0). + safe_hist = torch.minimum(cols, torch.full_like(cols, n_cols - 1)).expand( + rows, out_cols + ) + history_val = torch.gather(topk_result, 1, safe_hist) + + tail_raw = tail_start[:, None] + tail_off + tail_val = tail_raw.to(torch.int32) + if page_table is not None: + safe_tail = tail_raw.clamp(min=0, max=page_table.shape[1] - 1) + tail_val = torch.gather(page_table, 1, safe_tail).to(torch.int32) + elif topk_offsets is not None: + tail_val = (tail_raw + topk_offsets.to(torch.int64).unsqueeze(1)).to( + torch.int32 + ) + + out = torch.where(is_history, history_val, -1) + out = torch.where(is_tail, tail_val, out) + return out + + +@triton.jit +def _expand_pools_and_append_tail_kernel( + pool_ids_ptr, # [rows, n_groups], int (any int dtype) + seq_lens_ptr, # [rows], int32 (token-granular seq_len) + out_ptr, # [rows, out_cols], int32 + topk, # n_groups * pool_size + out_cols, # topk + pool_size - 1 + POOL_SIZE: tl.constexpr, + BLOCK_COLS: tl.constexpr, + pid_s0, + out_s0, +): + # Fuses expand_pools_to_tokens + append_tail_to_topk (identity path) into a + # single kernel. Each program writes one (row, column-tile) of the output. + row = tl.program_id(0) + tile = tl.program_id(1) + cols = tile * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + mask = cols < out_cols + + seq_len = tl.load(seq_lens_ptr + row) + pool_len = seq_len // POOL_SIZE + tail_start = pool_len * POOL_SIZE + tail_count = seq_len - tail_start # in [0, POOL_SIZE) + + # History region [0, topk): expand selected pool g = cols // POOL_SIZE. + is_history = cols < topk + g = cols // POOL_SIZE + o = cols % POOL_SIZE + pid = tl.load(pool_ids_ptr + row * pid_s0 + g, mask=mask & is_history, other=-1) + hist_val = (pid * POOL_SIZE + o).to(tl.int32) + hist_out = tl.where(pid >= 0, hist_val, -1) + + # Tail region [topk, out_cols): the request's trailing incomplete pool. + tail_off = cols - topk + is_tail = (tail_off >= 0) & (tail_off < tail_count) + tail_val = (tail_start + tail_off).to(tl.int32) + tail_out = tl.where(is_tail, tail_val, -1) + + result = tl.where(is_history, hist_out, tail_out) + tl.store(out_ptr + row * out_s0 + cols, result, mask=mask) + + +def expand_pools_and_append_tail( + pool_ids: torch.Tensor, + seq_lens: torch.Tensor, + pool_size: int, +) -> torch.Tensor: + """Fuse ``expand_pools_to_tokens`` + ``append_tail_to_topk`` (identity path). + + Produces the same ``[rows, topk + pool_size - 1]`` int32 output as calling + the two functions in sequence when neither ``page_table`` nor + ``topk_offsets`` is passed — the only path used by the GLM-5.3-Flash indexer. + The kernel derives ``pool_len = seq_len // pool_size`` internally, so the + caller no longer needs to precompute it. Replaces ~25 elementwise kernels + with one Triton launch. + """ + rows, n_groups = pool_ids.shape + topk = n_groups * pool_size + out_cols = topk + pool_size - 1 + out = torch.empty((rows, out_cols), dtype=torch.int32, device=pool_ids.device) + BLOCK_COLS = 128 + n_tiles = triton.cdiv(out_cols, BLOCK_COLS) + _expand_pools_and_append_tail_kernel[(rows, n_tiles)]( + pool_ids, + seq_lens, + out, + topk, + out_cols, + POOL_SIZE=pool_size, + BLOCK_COLS=BLOCK_COLS, + pid_s0=pool_ids.stride(0), + out_s0=out.stride(0), + ) + return out diff --git a/vllm/models/glm5next/sm70/__init__.py b/vllm/models/glm5next/sm70/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/glm5next/sm70/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/glm5next/sm70/fp8_kv.py b/vllm/models/glm5next/sm70/fp8_kv.py new file mode 100644 index 0000000000..4534f51274 --- /dev/null +++ b/vllm/models/glm5next/sm70/fp8_kv.py @@ -0,0 +1,381 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Packed E4M3FN KV-cache kernels for GLM-5.3 NoPE MLA on SM70.""" + +import torch + +from vllm.models.deepseek_v4.common.ops.fp8_software import ( + fp8_e4m3fn_bits_to_fp32_bitcast, + fp32_to_fp8_e4m3fn_bits, +) +from vllm.triton_utils import tl, triton + +GLM5_FP8_KV_DIM = 512 +GLM5_FP8_KV_GROUP_SIZE = 64 +GLM5_FP8_KV_NUM_SCALES = GLM5_FP8_KV_DIM // GLM5_FP8_KV_GROUP_SIZE +GLM5_FP8_KV_SLOT_BYTES = GLM5_FP8_KV_DIM + GLM5_FP8_KV_NUM_SCALES + + +@triton.jit +def _sm70_glm5_fp8_kv_insert_kernel( + kv_ptr, + slot_mapping_ptr, + cache_ptr, + cache_stride_block, + block_size: tl.constexpr, + HEAD_DIM: tl.constexpr, + GROUP_SIZE: tl.constexpr, + NUM_GROUPS: tl.constexpr, +): + token_idx = tl.program_id(0) + slot_idx = tl.load(slot_mapping_ptr + token_idx) + if slot_idx < 0: + return + + offsets = tl.arange(0, HEAD_DIM) + values = tl.load(kv_ptr + token_idx * HEAD_DIM + offsets).to(tl.float32) + grouped = tl.reshape(values, (NUM_GROUPS, GROUP_SIZE)) + absmax = tl.max(tl.abs(grouped), axis=1) + absmax = tl.maximum(absmax, 1.0e-4) + exponent = tl.ceil(tl.log2(absmax / 448.0)) + scales = tl.exp2(exponent) + quantized = tl.clamp(grouped / scales[:, None], -448.0, 448.0) + packed = fp32_to_fp8_e4m3fn_bits(quantized) + + block_idx = slot_idx // block_size + pos_in_block = slot_idx % block_size + block_base = cache_ptr + block_idx.to(tl.int64) * cache_stride_block + token_data = block_base + pos_in_block * HEAD_DIM + tl.store(token_data + offsets, tl.reshape(packed, (HEAD_DIM,))) + + scale_offsets = tl.arange(0, NUM_GROUPS) + encoded_scales = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.uint8) + token_scales = block_base + block_size * HEAD_DIM + pos_in_block * NUM_GROUPS + tl.store(token_scales + scale_offsets, encoded_scales) + + +@triton.jit +def _sm70_glm5_sparse_paged_fp8_kernel( + q_ptr, + cache_ptr, + indices_ptr, + lengths_ptr, + out_ptr, + q_stride_t, + q_stride_h, + cache_stride_block, + indices_stride_t, + out_stride_t, + out_stride_h, + num_heads, + num_cache_slots, + block_size, + scale, + INDEX_WIDTH: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + HEAD_DIM: tl.constexpr, + GROUP_SIZE: tl.constexpr, +): + query_idx = tl.program_id(0) + head_offsets = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) + head_mask = head_offsets < num_heads + dim_offsets = tl.arange(0, HEAD_DIM) + + q = tl.load( + q_ptr + + query_idx * q_stride_t + + head_offsets[:, None] * q_stride_h + + dim_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) + + neg_large = -3.4028234663852886e38 + running_max = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + running_sum = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc = tl.zeros((BLOCK_H, HEAD_DIM), dtype=tl.float32) + # Keep the loop bound runtime-visible so CUDA graph replay only scans the + # selected keys for the current sequence. INDEX_WIDTH is the fixed graph + # buffer capacity (2048 for GLM-5.3), not the amount of useful work. + valid_len = tl.minimum(tl.load(lengths_ptr + query_idx), INDEX_WIDTH) + key_offsets = tl.arange(0, BLOCK_K) + + for start in tl.range(0, valid_len, BLOCK_K): + positions = start + key_offsets + slots = tl.load( + indices_ptr + query_idx * indices_stride_t + positions, + mask=positions < INDEX_WIDTH, + other=-1, + ) + valid = (positions < valid_len) & (slots >= 0) & (slots < num_cache_slots) + safe_slots = tl.where(valid, slots, 0) + block_idx = safe_slots // block_size + pos_in_block = safe_slots % block_size + block_base = cache_ptr + block_idx.to(tl.int64) * cache_stride_block + token_data = block_base + pos_in_block * HEAD_DIM + token_scales = ( + block_base + block_size * HEAD_DIM + pos_in_block * (HEAD_DIM // GROUP_SIZE) + ) + + packed = tl.load( + token_data[:, None] + dim_offsets[None, :], + mask=valid[:, None], + other=0, + ) + fp8 = fp8_e4m3fn_bits_to_fp32_bitcast(packed) + scale_group_offsets = tl.arange(0, HEAD_DIM // GROUP_SIZE) + encoded_scale = tl.load( + token_scales[:, None] + scale_group_offsets[None, :], + mask=valid[:, None], + other=127, + ) + dequant_scale = tl.exp2(encoded_scale.to(tl.float32) - 127.0) + fp8_grouped = tl.reshape(fp8, (BLOCK_K, HEAD_DIM // GROUP_SIZE, GROUP_SIZE)) + scale_grouped = tl.reshape( + dequant_scale.to(tl.float16), + (BLOCK_K, HEAD_DIM // GROUP_SIZE, 1), + ) + kv = tl.reshape( + fp8_grouped.to(tl.float16) * scale_grouped, + (BLOCK_K, HEAD_DIM), + ) + kv = tl.where(valid[:, None], kv, 0.0) + + scores = tl.dot(q, tl.trans(kv)) * scale + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) + block_max = tl.max(scores, axis=1) + new_max = tl.maximum(running_max, block_max) + alpha = tl.exp(running_max - new_max) + probs = tl.exp(scores - new_max[:, None]) + probs = tl.where(head_mask[:, None] & valid[None, :], probs, 0.0) + acc = acc * alpha[:, None] + tl.dot(probs.to(tl.float16), kv) + running_sum = running_sum * alpha + tl.sum(probs, axis=1) + running_max = new_max + + denom = tl.maximum(running_sum, 1.0e-30) + result = tl.where(running_sum[:, None] > 0.0, acc / denom[:, None], 0.0) + tl.store( + out_ptr + + query_idx * out_stride_t + + head_offsets[:, None] * out_stride_h + + dim_offsets[None, :], + result, + mask=head_mask[:, None], + ) + + +@triton.jit +def _sm70_glm5_fp8_kv_gather_dequant_kernel( + cache_ptr, + indices_ptr, + lengths_ptr, + gathered_ptr, + cache_stride_block, + num_cache_slots, + block_size, + INDEX_WIDTH: tl.constexpr, + HEAD_DIM: tl.constexpr, + GROUP_SIZE: tl.constexpr, +): + position = tl.program_id(0) + dim_offsets = tl.arange(0, HEAD_DIM) + valid_len = tl.minimum(tl.load(lengths_ptr), INDEX_WIDTH) + slot = tl.load(indices_ptr + position, mask=position < INDEX_WIDTH, other=-1) + valid = (position < valid_len) & (slot >= 0) & (slot < num_cache_slots) + safe_slot = tl.where(valid, slot, 0) + block_idx = safe_slot // block_size + pos_in_block = safe_slot % block_size + block_base = cache_ptr + block_idx.to(tl.int64) * cache_stride_block + token_data = block_base + pos_in_block * HEAD_DIM + token_scales = ( + block_base + block_size * HEAD_DIM + pos_in_block * (HEAD_DIM // GROUP_SIZE) + ) + + packed = tl.load(token_data + dim_offsets, mask=valid, other=0) + fp8 = fp8_e4m3fn_bits_to_fp32_bitcast(packed) + scale_group_offsets = tl.arange(0, HEAD_DIM // GROUP_SIZE) + encoded_scale = tl.load( + token_scales + scale_group_offsets, + mask=valid, + other=127, + ) + scales = tl.exp2(encoded_scale.to(tl.float32) - 127.0).to(tl.float16) + grouped = tl.reshape(fp8, (HEAD_DIM // GROUP_SIZE, GROUP_SIZE)) + dequantized = tl.reshape( + grouped.to(tl.float16) * scales[:, None], + (HEAD_DIM,), + ) + tl.store( + gathered_ptr + position * HEAD_DIM + dim_offsets, + tl.where(valid, dequantized, 0.0), + ) + + +@triton.jit +def _sm70_glm5_sparse_scores_softmax_kernel( + scores_ptr, + indices_ptr, + lengths_ptr, + probs_ptr, + scores_stride_h, + num_cache_slots, + scale, + INDEX_WIDTH: tl.constexpr, + WIDTH_BLOCK: tl.constexpr, +): + head_idx = tl.program_id(0) + positions = tl.arange(0, WIDTH_BLOCK) + width_mask = positions < INDEX_WIDTH + valid_len = tl.minimum(tl.load(lengths_ptr), INDEX_WIDTH) + slots = tl.load(indices_ptr + positions, mask=width_mask, other=-1) + valid = ( + width_mask & (positions < valid_len) & (slots >= 0) & (slots < num_cache_slots) + ) + neg_large = -3.4028234663852886e38 + scores = tl.load( + scores_ptr + head_idx * scores_stride_h + positions, + mask=width_mask, + other=neg_large, + ).to(tl.float32) + scores = tl.where(valid, scores * scale, neg_large) + scores_max = tl.max(scores, axis=0) + numerators = tl.exp(scores - scores_max) + numerators = tl.where(valid, numerators, 0.0) + denominator = tl.sum(numerators, axis=0) + probs = tl.where( + denominator > 0.0, + numerators / tl.maximum(denominator, 1.0e-30), + 0.0, + ) + tl.store( + probs_ptr + head_idx * scores_stride_h + positions, + probs, + mask=width_mask, + ) + + +def sm70_glm5_fp8_kv_insert( + kv: torch.Tensor, + cache: torch.Tensor, + slot_mapping: torch.Tensor, +) -> None: + """Quantize FP16 latent KV into the packed GLM-5.3 paged layout.""" + assert kv.dtype == torch.float16 and kv.ndim == 2 + assert kv.shape[1] == GLM5_FP8_KV_DIM and kv.is_contiguous() + assert cache.dtype == torch.uint8 and cache.ndim == 3 + assert cache.shape[2] == GLM5_FP8_KV_SLOT_BYTES + num_tokens = slot_mapping.numel() + if num_tokens == 0: + return + _sm70_glm5_fp8_kv_insert_kernel[(num_tokens,)]( + kv, + slot_mapping, + cache, + cache.stride(0), + block_size=cache.shape[1], + HEAD_DIM=GLM5_FP8_KV_DIM, + GROUP_SIZE=GLM5_FP8_KV_GROUP_SIZE, + NUM_GROUPS=GLM5_FP8_KV_NUM_SCALES, + num_warps=8, + ) + + +def sm70_glm5_sparse_attention_paged_fp8( + q: torch.Tensor, + cache: torch.Tensor, + indices: torch.Tensor, + lengths: torch.Tensor, + scale: float, + out: torch.Tensor, +) -> None: + """Run sparse NoPE MLA directly over the packed GLM-5.3 FP8 cache.""" + assert q.dtype == out.dtype == torch.float16 + assert q.shape == out.shape and q.shape[-1] == GLM5_FP8_KV_DIM + assert cache.dtype == torch.uint8 and cache.ndim == 3 + assert cache.shape[2] == GLM5_FP8_KV_SLOT_BYTES + indices_2d = indices.reshape(q.shape[0], -1) + lengths_1d = lengths.reshape(-1).to(torch.int32) + assert indices_2d.shape[0] == q.shape[0] == lengths_1d.shape[0] + + block_h = 8 + _sm70_glm5_sparse_paged_fp8_kernel[(q.shape[0], triton.cdiv(q.shape[1], block_h))]( + q, + cache, + indices_2d, + lengths_1d, + out, + q.stride(0), + q.stride(1), + cache.stride(0), + indices_2d.stride(0), + out.stride(0), + out.stride(1), + q.shape[1], + cache.shape[0] * cache.shape[1], + cache.shape[1], + float(scale), + INDEX_WIDTH=indices_2d.shape[1], + BLOCK_H=block_h, + BLOCK_K=16, + HEAD_DIM=GLM5_FP8_KV_DIM, + GROUP_SIZE=GLM5_FP8_KV_GROUP_SIZE, + num_warps=4, + ) + + +def sm70_glm5_sparse_attention_paged_fp8_gemm( + q: torch.Tensor, + cache: torch.Tensor, + indices: torch.Tensor, + lengths: torch.Tensor, + scale: float, + out: torch.Tensor, + gathered_kv: torch.Tensor, + scores: torch.Tensor, + probs: torch.Tensor, +) -> None: + """Run B1 sparse MLA with an SM70 dequant + tensor-core GEMM pipeline.""" + assert q.dtype == out.dtype == torch.float16 + assert q.shape == out.shape and q.shape[-1] == GLM5_FP8_KV_DIM + assert q.shape[0] == 1 + assert cache.dtype == torch.uint8 and cache.ndim == 3 + assert cache.shape[2] == GLM5_FP8_KV_SLOT_BYTES + indices_2d = indices.reshape(1, -1) + lengths_1d = lengths.reshape(-1).to(torch.int32) + index_width = indices_2d.shape[1] + assert lengths_1d.shape == (1,) + assert gathered_kv.shape == (index_width, GLM5_FP8_KV_DIM) + assert gathered_kv.dtype == torch.float16 and gathered_kv.is_contiguous() + assert scores.shape == probs.shape == (q.shape[1], index_width) + assert scores.dtype == probs.dtype == torch.float16 + assert scores.is_contiguous() and probs.is_contiguous() + + _sm70_glm5_fp8_kv_gather_dequant_kernel[(index_width,)]( + cache, + indices_2d, + lengths_1d, + gathered_kv, + cache.stride(0), + cache.shape[0] * cache.shape[1], + cache.shape[1], + INDEX_WIDTH=index_width, + HEAD_DIM=GLM5_FP8_KV_DIM, + GROUP_SIZE=GLM5_FP8_KV_GROUP_SIZE, + num_warps=4, + ) + torch.mm(q[0], gathered_kv.t(), out=scores) + _sm70_glm5_sparse_scores_softmax_kernel[(q.shape[1],)]( + scores, + indices_2d, + lengths_1d, + probs, + scores.stride(0), + cache.shape[0] * cache.shape[1], + float(scale), + INDEX_WIDTH=index_width, + WIDTH_BLOCK=triton.next_power_of_2(index_width), + num_warps=8, + ) + torch.mm(probs, gathered_kv, out=out[0]) diff --git a/vllm/models/glm5next/sm70/sparse.py b/vllm/models/glm5next/sm70/sparse.py new file mode 100644 index 0000000000..90a5d64a53 --- /dev/null +++ b/vllm/models/glm5next/sm70/sparse.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""GLM-5.3-Flash sparse MLA backend for exact SM70 CUDA devices.""" + +from typing import TYPE_CHECKING, ClassVar + +import torch + +from vllm.config import get_current_vllm_config_or_none +from vllm.config.cache import CacheDType +from vllm.logger import init_logger +from vllm.models.deepseek_v4.sm70.sparse_kernels import ( + sm70_sparse_attention_gathered, +) +from vllm.models.glm5next.sm70.fp8_kv import ( + GLM5_FP8_KV_SLOT_BYTES, + sm70_glm5_fp8_kv_insert, + sm70_glm5_sparse_attention_paged_fp8, + sm70_glm5_sparse_attention_paged_fp8_gemm, +) +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backend import AttentionLayer, SparseMLAAttentionImpl +from vllm.v1.attention.backends.mla.flashmla_sparse import ( + FlashMLASparseBackend, + FlashMLASparseMetadata, +) +from vllm.v1.attention.backends.mla.sparse_utils import ( + triton_convert_req_index_to_global_index, +) +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.model_executor.models.deepseek_v2 import Indexer + +logger = init_logger(__name__) + + +class Glm5NextSM70SparseBackend(FlashMLASparseBackend): + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "float16", + "fp8", + "fp8_e4m3", + ] + + @staticmethod + def get_name() -> str: + return "GLM5_SM70_SPARSE" + + @staticmethod + def get_impl_cls() -> type["Glm5NextSM70SparseImpl"]: + return Glm5NextSM70SparseImpl + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [512] + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + del num_kv_heads + if cache_dtype_str in {"fp8", "fp8_e4m3"}: + return (num_blocks, block_size, GLM5_FP8_KV_SLOT_BYTES) + return (num_blocks, block_size, head_size) + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + return capability.major == 7 and capability.minor == 0 + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: str | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + device_capability: DeviceCapability, + ) -> str | None: + del head_size, dtype, kv_cache_dtype, block_size, use_mla, has_sink + del use_sparse, device_capability + config = get_current_vllm_config_or_none() + if config is None or config.model_config is None: + return None + text_config = config.model_config.hf_text_config + if ( + getattr(text_config, "qk_rope_head_dim", None) != 0 + or getattr(text_config, "kv_lora_rank", None) != 512 + ): + return "GLM5_SM70_SPARSE requires NoPE MLA with kv_lora_rank=512" + return None + + +class Glm5NextSM70SparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]): + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None, + attn_type: str, + kv_sharing_target_layer_name: str | None, + topk_indices_buffer: torch.Tensor | None = None, + indexer: "Indexer | None" = None, + **mla_args, + ) -> None: + del alibi_slopes, sliding_window, logits_soft_cap, attn_type + del kv_sharing_target_layer_name, topk_indices_buffer + if num_kv_heads != 1 or head_size != 512: + raise NotImplementedError( + "GLM-5.3 SM70 sparse MLA requires one 512-wide latent KV head." + ) + if kv_cache_dtype not in {"auto", "float16", "fp8", "fp8_e4m3"}: + raise NotImplementedError( + "GLM-5.3 SM70 sparse MLA supports FP16 or E4M3 KV cache." + ) + if mla_args["qk_rope_head_dim"] != 0 or mla_args["kv_lora_rank"] != 512: + raise NotImplementedError( + "GLM-5.3 SM70 sparse MLA requires NoPE and kv_lora_rank=512." + ) + if indexer is None: + raise ValueError("GLM-5.3 sparse MLA requires its KPool indexer.") + + self.num_heads = num_heads + self.head_size = head_size + self.scale = float(scale) + self.softmax_scale = float(scale) + self.num_kv_heads = num_kv_heads + self.kv_cache_dtype = kv_cache_dtype + self.use_fp8_cache = kv_cache_dtype in {"fp8", "fp8_e4m3"} + self.kv_lora_rank = 512 + topk_indices_buffer = indexer.topk_indices_buffer + if topk_indices_buffer is None: + raise ValueError("GLM-5.3 sparse MLA requires an index buffer.") + self.topk_indices_buffer = topk_indices_buffer + self.index_width = self.topk_indices_buffer.shape[1] + + config = get_current_vllm_config_or_none() + if config is None: + raise RuntimeError("GLM-5.3 SM70 sparse MLA requires VllmConfig.") + max_tokens = config.scheduler_config.max_num_batched_tokens + workspace_specs: list[tuple[tuple[int, ...], torch.dtype]] = [ + ((max_tokens, num_heads, self.kv_lora_rank), torch.float16), + ] + if self.use_fp8_cache: + workspace_specs.extend( + ( + ((self.index_width, self.kv_lora_rank), torch.float16), + ((num_heads, self.index_width), torch.float16), + ((num_heads, self.index_width), torch.float16), + ) + ) + current_workspace_manager().get_simultaneous(*workspace_specs) + logger.info_once( + "GLM-5.3-Flash route: SM70 FP16 sparse MLA with %s KV%s.", + "packed E4M3FN" if self.use_fp8_cache else "FP16", + (" and B1 dequant + tensor-core GEMM decode" if self.use_fp8_cache else ""), + ) + + def do_kv_cache_update( + self, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + ) -> None: + if not self.use_fp8_cache: + return super().do_kv_cache_update( + kv_c_normed, + k_pe, + kv_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + ) + del k_scale + if kv_cache.numel() == 0 or slot_mapping is None: + return + if k_pe.shape[-1] != 0: + raise NotImplementedError("GLM-5.3 packed FP8 KV is NoPE-only.") + sm70_glm5_fp8_kv_insert( + kv_c_normed, + kv_cache.view(torch.uint8), + slot_mapping.flatten(), + ) + + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: FlashMLASparseMetadata, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + del layer + if isinstance(q, tuple): + q_nope, q_pe = q + if q_pe.shape[-1] != 0: + raise NotImplementedError("GLM-5.3 SM70 sparse MLA is NoPE-only.") + q = q_nope + + if q.dtype != torch.float16: + raise TypeError("GLM-5.3 SM70 sparse MLA requires an FP16 query.") + if q.shape[-1] != self.kv_lora_rank: + raise ValueError(f"Expected latent query width 512, got {q.shape[-1]}.") + + num_tokens = q.shape[0] + topk_indices = self.topk_indices_buffer[:num_tokens] + global_indices, valid_counts = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token, + attn_metadata.block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + return_valid_counts=True, + ) + workspace_manager = current_workspace_manager() + if self.use_fp8_cache: + if num_tokens == 1: + out, gathered_kv, scores, probs = workspace_manager.get_simultaneous( + ((1, self.num_heads, self.kv_lora_rank), torch.float16), + ((self.index_width, self.kv_lora_rank), torch.float16), + ((self.num_heads, self.index_width), torch.float16), + ((self.num_heads, self.index_width), torch.float16), + ) + sm70_glm5_sparse_attention_paged_fp8_gemm( + q, + kv_c_and_k_pe_cache.view(torch.uint8), + global_indices, + valid_counts, + self.softmax_scale, + out, + gathered_kv, + scores, + probs, + ) + else: + (out,) = workspace_manager.get_simultaneous( + ( + (num_tokens, self.num_heads, self.kv_lora_rank), + torch.float16, + ), + ) + sm70_glm5_sparse_attention_paged_fp8( + q, + kv_c_and_k_pe_cache.view(torch.uint8), + global_indices, + valid_counts, + self.softmax_scale, + out, + ) + else: + if kv_c_and_k_pe_cache.dtype != torch.float16: + raise TypeError("GLM-5.3 FP16 KV route received a non-FP16 cache.") + (out,) = workspace_manager.get_simultaneous( + ( + (num_tokens, self.num_heads, self.kv_lora_rank), + torch.float16, + ), + ) + sm70_sparse_attention_gathered( + q, + kv_c_and_k_pe_cache, + global_indices, + valid_counts, + self.softmax_scale, + None, + out, + ) + return out, None diff --git a/vllm/models/qwen4_exp/__init__.py b/vllm/models/qwen4_exp/__init__.py new file mode 100644 index 0000000000..9e533e28b4 --- /dev/null +++ b/vllm/models/qwen4_exp/__init__.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp model package.""" + +from typing import TYPE_CHECKING, Any + +from .common.hyperconnection import ( + GatedResidual, + GroupedGemmaRMSNorm, + HyperConnectionBase, + HyperConnectionConfig, +) + +if TYPE_CHECKING: + from .nvidia.model import ( + Qwen4ExpForCausalLM, + Qwen4ExpForConditionalGeneration, + ) + from .nvidia.mtp import Qwen4ExpMTP + + +def __getattr__(name: str) -> Any: + if name in { + "Qwen4ExpForCausalLM", + "Qwen4ExpForConditionalGeneration", + "Qwen4ExpMTP", + }: + from vllm.platforms import current_platform + + if current_platform.is_xpu() or current_platform.is_tpu(): + raise NotImplementedError("Qwen4Exp currently supports CUDA and ROCm only") + if current_platform.is_rocm(): + from .amd.model import ( + Qwen4ExpForCausalLM, + Qwen4ExpForConditionalGeneration, + ) + from .amd.mtp import Qwen4ExpMTP + else: + from .nvidia.model import ( + Qwen4ExpForCausalLM, + Qwen4ExpForConditionalGeneration, + ) + from .nvidia.mtp import Qwen4ExpMTP + + return { + "Qwen4ExpForCausalLM": Qwen4ExpForCausalLM, + "Qwen4ExpForConditionalGeneration": (Qwen4ExpForConditionalGeneration), + "Qwen4ExpMTP": Qwen4ExpMTP, + }[name] + raise AttributeError(name) + + +__all__ = [ + "GatedResidual", + "GroupedGemmaRMSNorm", + "HyperConnectionBase", + "HyperConnectionConfig", + "Qwen4ExpForCausalLM", + "Qwen4ExpForConditionalGeneration", + "Qwen4ExpMTP", +] diff --git a/vllm/models/qwen4_exp/common/__init__.py b/vllm/models/qwen4_exp/common/__init__.py new file mode 100644 index 0000000000..a89d77b976 --- /dev/null +++ b/vllm/models/qwen4_exp/common/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Common Qwen4Exp model components.""" + +from .hyperconnection import ( + GatedResidual, + GroupedGemmaRMSNorm, + HyperConnectionBase, + HyperConnectionConfig, +) + +__all__ = [ + "GatedResidual", + "GroupedGemmaRMSNorm", + "HyperConnectionBase", + "HyperConnectionConfig", +] diff --git a/vllm/models/qwen4_exp/common/hyperconnection.py b/vllm/models/qwen4_exp/common/hyperconnection.py new file mode 100644 index 0000000000..5f8bde0479 --- /dev/null +++ b/vllm/models/qwen4_exp/common/hyperconnection.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HyperConnection (Gated Residual) utilities. + +Implements the HyperConnection residual scheme proposed in +"HyperConnections" (https://arxiv.org/abs/2409.19606). + +The two concrete variants are: + - ``HyperConnectionBase`` - simple average pooling across hc_count parallel + streams (equivalent to hyperconnection_average). + - ``GatedResidual`` - learnable low-rank gated mixing and injection + (gated_residual). + +Hidden states between layers have shape ``[..., HC*HS]`` with HS inner +(HC outer, HS inner — checkpoint-native layout). The local torch +implementation consumes the hyper input viewed as ``[..., HC, HS]``. + +Typical usage inside a transformer decoder layer:: + + self.attn_hc = GatedResidual(hc_config, role="attn") + self.mlp_hc = GatedResidual(hc_config, role="mlp") + + hidden_states, residual = self.attn_hc.mix(hidden_states) + hidden_states = attention(hidden_states) + hidden_states = self.attn_hc.combine(hidden_states, residual) + + hidden_states, residual = self.mlp_hc.mix(hidden_states) + hidden_states = mlp(hidden_states) + hidden_states = self.mlp_hc.combine(hidden_states, residual) +""" + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch import nn + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- +@dataclass +class HyperConnectionConfig: + """Configuration shared by all HyperConnection variants.""" + + hc_count: int = 4 + hidden_size: int = 64 + params_dtype: torch.dtype = torch.bfloat16 + mtp_hc: bool = False + hc_lowrank: int = 16 + rms_norm_eps: float = 1e-6 + hc_per_branch_norm: bool = False + + +class GroupedGemmaRMSNorm(nn.Module): + def __init__( + self, + hidden_size: int, + eps: float, + group_size: int | None, + dtype: torch.dtype | None, + ) -> None: + super().__init__() + if group_size is not None and hidden_size % group_size: + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by " + f"group_size ({group_size})" + ) + self.variance_epsilon = eps + self.group_size = group_size + self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=dtype)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.float() + if self.group_size is None: + variance = hidden_states.square().mean(dim=-1, keepdim=True) + normalized = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + else: + grouped = hidden_states.unflatten( + -1, (hidden_states.shape[-1] // self.group_size, self.group_size) + ) + variance = grouped.square().mean(dim=-1, keepdim=True) + normalized = ( + grouped * torch.rsqrt(variance + self.variance_epsilon) + ).flatten(-2) + return (normalized * (1.0 + self.weight.float())).to(input_dtype) + + +# --------------------------------------------------------------------------- +# Average-pooling variant +# --------------------------------------------------------------------------- +class HyperConnectionBase(nn.Module): + """Average-pooling HyperConnection (``hyperconnection_average``). + + Splits the incoming ``[..., HC*HS]`` tensor (HC outer, HS inner) into + ``HC`` parallel streams, averages them for the block input, and + broadcasts the block output back to every stream. + """ + + def __init__( + self, + config: HyperConnectionConfig, + layer_idx: int | None = None, + role: str | None = None, + ) -> None: + super().__init__() + self.config = config + self.hc_count = config.hc_count + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + self.role = role + + @property + def hyper_hidden_size(self) -> int: + return self.hc_count * self.hidden_size + + def mix(self, hyper_input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Average the HC streams into a single block input.""" + assert hyper_input.shape[-1] == self.hc_count * self.hidden_size + # [*, HC, HS] — mean over HC (dim=-2). + unflat = hyper_input.unflatten(-1, (self.hc_count, self.hidden_size)) + mixed_input = unflat.mean(dim=-2) + return mixed_input, hyper_input + + def combine( + self, block_output: torch.Tensor, residual: torch.Tensor + ) -> torch.Tensor: + """Broadcast the block output back to every stream.""" + assert residual.shape[-1] == self.hc_count * self.hidden_size + assert block_output.shape[-1] == self.hidden_size + residual_reshaped = residual.unflatten(-1, (self.hc_count, self.hidden_size)) + combined = residual_reshaped + block_output.unsqueeze(-2) + return combined.flatten(-2) + + +# --------------------------------------------------------------------------- +# Gated-residual variant +# --------------------------------------------------------------------------- +class GatedResidual(HyperConnectionBase): + """Gated HyperConnection with learnable low-rank mixing and injection. + + ``mix()`` applies GemmaRMSNorm per HC stream and projects through a + low-rank sigmoid gate to produce a single block input. ``combine()`` + injects the block output back into each stream through a learned + per-stream injection weight. + + This implementation uses only PyTorch operators. Tensor-parallel + collectives are supplied by its caller. + """ + + def __init__( + self, + config: HyperConnectionConfig, + layer_idx: int | None = None, + role: str | None = None, + use_mix: bool = True, + use_combine: bool = True, + ) -> None: + super().__init__(config, layer_idx, role) + norm_size = ( + self.hyper_hidden_size if config.hc_per_branch_norm else config.hidden_size + ) + group_size = config.hidden_size if config.hc_per_branch_norm else None + # Normalize each H-sized HC stream independently while retaining a + # separate affine weight for every element of the HC*H layout. + self.hc_norm = GroupedGemmaRMSNorm( + norm_size, + eps=config.rms_norm_eps, + group_size=group_size, + dtype=config.params_dtype, + ) + + # -- raw Linear weights (checkpoint-compatible) ---------------------- + if use_mix: + self.input_mix_weight_down = nn.Linear( + self.hyper_hidden_size, + config.hc_lowrank, + bias=False, + dtype=config.params_dtype, + ) + self.input_mix_weight_up = nn.Linear( + config.hc_lowrank, + self.hyper_hidden_size, + bias=False, + dtype=config.params_dtype, + ) + if use_combine: + self.block_inject_weight = nn.Linear( + self.hyper_hidden_size, + self.hc_count, + bias=False, + dtype=config.params_dtype, + ) + + def _normalize(self, hyper_input: torch.Tensor) -> torch.Tensor: + if self.config.hc_per_branch_norm: + return self.hc_norm(hyper_input) + return self.hc_norm( + hyper_input.unflatten(-1, (self.hc_count, self.hidden_size)) + ).flatten(-2) + + def mix( + self, hyper_input: torch.Tensor + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """Mix: RMSNorm -> low-rank gate -> gated mean.""" + assert hyper_input.shape[-1] == self.hc_count * self.hidden_size + if not hasattr(self, "input_mix_weight_down"): + raise RuntimeError("mix was disabled for this hyper-connection") + hyper_input_normed = self._normalize(hyper_input) + # Gate — original mix order: linear+silu then linear+sigmoid. + gate = F.silu( + F.linear(hyper_input_normed, self.input_mix_weight_down.weight) + / self.hc_count + ) + gate = torch.sigmoid(F.linear(gate, self.input_mix_weight_up.weight)).unflatten( + -1, (self.hc_count, self.hidden_size) + ) + mixed_input = ( + gate * hyper_input_normed.unflatten(-1, (self.hc_count, self.hidden_size)) + ).mean(dim=-2) + return mixed_input.to(hyper_input.dtype), (hyper_input, hyper_input_normed) + + def combine( + self, + block_output: torch.Tensor, + residuals: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + if not hasattr(self, "block_inject_weight"): + raise RuntimeError("combine was disabled for this hyper-connection") + hyper_input, hyper_input_normed = residuals + assert hyper_input.shape[-1] == self.hc_count * self.hidden_size + assert block_output.shape[-1] == self.hidden_size + residual = hyper_input.unflatten(-1, (self.hc_count, self.hidden_size)) + # The paired mix keeps its normalized hyper input so combine uses the + # same HC module's injection weight. + injection_weight = 2.0 * torch.sigmoid( + F.linear(hyper_input_normed, self.block_inject_weight.weight) + / self.hc_count + ) + output = residual + block_output.unsqueeze(-2) * injection_weight.unsqueeze(-1) + return output.flatten(-2).to(hyper_input.dtype) + + +__all__ = [ + "GatedResidual", + "GroupedGemmaRMSNorm", + "HyperConnectionBase", + "HyperConnectionConfig", +] diff --git a/vllm/models/qwen4_exp/common/ple.py b/vllm/models/qwen4_exp/common/ple.py new file mode 100644 index 0000000000..8047f1ad10 --- /dev/null +++ b/vllm/models/qwen4_exp/common/ple.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Common Qwen4Exp PLE helpers.""" + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class PLEShardOverlap: + """Source and destination slices for one checkpoint embedding shard.""" + + source_start: int + destination_start: int + row_count: int + + +def compute_ple_shard_overlap( + *, + checkpoint_start: int, + checkpoint_rows: int, + tp_start: int, + tp_end: int, +) -> PLEShardOverlap | None: + """Compute the overlap of a checkpoint shard and one TP vocabulary range.""" + + if checkpoint_start < 0 or checkpoint_rows < 0: + raise ValueError("checkpoint shard bounds must be non-negative") + if tp_start < 0 or tp_end < tp_start: + raise ValueError("invalid TP vocabulary range") + checkpoint_end = checkpoint_start + checkpoint_rows + overlap_start = max(checkpoint_start, tp_start) + overlap_end = min(checkpoint_end, tp_end) + if overlap_start >= overlap_end: + return None + return PLEShardOverlap( + source_start=overlap_start - checkpoint_start, + destination_start=overlap_start - tp_start, + row_count=overlap_end - overlap_start, + ) + + +def copy_ple_embedding_shard_( + destination: torch.Tensor, + loaded_weight: torch.Tensor, + *, + checkpoint_start: int, + tp_start: int, + tp_end: int, +) -> int: + """Copy the overlapping rows of a PLE checkpoint shard into a TP table.""" + + if destination.ndim == 0 or loaded_weight.ndim != destination.ndim: + raise ValueError("destination and loaded weight must have matching ranks") + if destination.shape[1:] != loaded_weight.shape[1:]: + raise ValueError( + "embedding shard dimensions do not match: " + f"{tuple(destination.shape[1:])} != {tuple(loaded_weight.shape[1:])}" + ) + if destination.shape[0] < tp_end - tp_start: + raise ValueError("destination does not cover the requested TP range") + overlap = compute_ple_shard_overlap( + checkpoint_start=checkpoint_start, + checkpoint_rows=loaded_weight.shape[0], + tp_start=tp_start, + tp_end=tp_end, + ) + if overlap is None: + return 0 + source = loaded_weight.narrow(0, overlap.source_start, overlap.row_count) + target = destination.narrow(0, overlap.destination_start, overlap.row_count) + with torch.no_grad(): + target.copy_(source.to(device=target.device, dtype=target.dtype)) + return overlap.row_count diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py new file mode 100644 index 0000000000..bb5f890cef --- /dev/null +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -0,0 +1,824 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Paged side-cache ownership and metadata for Qwen4Exp QSA. + +Each QSA layer keeps a fixed circular buffer of raw index keys (the +compressor state) and one compressed key. MRoPE models pack exact three-axis +positions beside the raw keys; text models derive group positions from +logical positions. The compressor state uses one block per request, while +the compressed owner uses ``MLAAttentionSpec.compress_ratio`` so its block +table follows the main KV-cache lifecycle. Their physical tensor storage is +shared by the generic cache-layout planner. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from functools import cache +from typing import ClassVar + +import torch +from torch import nn + +from vllm.config import CacheConfig, VllmConfig +from vllm.config.cache import CacheDType +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.utils.math_utils import cdiv +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadata, + AttentionMetadataBuilder, + CommonAttentionMetadata, +) +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + CircularBufferSpec, + KVCacheSpec, + MLAAttentionSpec, +) + + +def canonical_qsa_rope_positions(positions: torch.Tensor) -> torch.Tensor: + """Return exact per-token positions as ``[tokens, 1, 3]`` int64 rows.""" + + if positions.ndim == 1: + positions = positions.unsqueeze(0).expand(3, -1) + elif positions.ndim != 2 or positions.shape[0] not in (1, 3): + raise ValueError("QSA RoPE positions must be [tokens] or [1|3, tokens]") + if positions.shape[0] == 1: + positions = positions.expand(3, -1) + return positions.transpose(0, 1).unsqueeze(1).to(torch.int64) + + +def _logical_positions( + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + token_to_req: torch.Tensor, + num_tokens: int, +) -> torch.Tensor: + if num_tokens == 0: + return seq_lens.new_empty((0,), dtype=torch.int64) + arange = torch.arange(num_tokens, device=query_start_loc.device) + requests = token_to_req[:num_tokens].long() + query_lens = torch.diff(query_start_loc) + within_query = arange - query_start_loc.index_select(0, requests) + return ( + seq_lens.index_select(0, requests).long() + - query_lens.index_select(0, requests).long() + + within_query.long() + ) + + +def _logical_to_physical_qsa_slots( + block_table: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_size: int, +) -> torch.Tensor: + if block_size <= 0: + raise ValueError("QSA cache block size must be positive") + if block_table.ndim != 2: + raise ValueError("QSA block table must be two-dimensional") + if request_indices.shape != logical_positions.shape: + request_indices = torch.broadcast_to(request_indices, logical_positions.shape) + + requests = request_indices.to(device=block_table.device, dtype=torch.long) + positions = logical_positions.to(device=block_table.device, dtype=torch.long) + valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0) + logical_blocks = torch.div( + positions.clamp_min(0), block_size, rounding_mode="floor" + ) + valid &= logical_blocks < block_table.shape[1] + safe_requests = requests.clamp(0, max(block_table.shape[0] - 1, 0)) + safe_blocks = logical_blocks.clamp(0, max(block_table.shape[1] - 1, 0)) + if not all(block_table.shape): + return torch.full_like(positions, PAD_SLOT_ID) + physical_blocks = block_table[safe_requests, safe_blocks].long() + valid &= physical_blocks >= 0 + slots = physical_blocks * block_size + positions.remainder(block_size) + return torch.where(valid, slots, PAD_SLOT_ID) + + +def circular_qsa_slot_mapping( + block_table: torch.Tensor, + token_to_req: torch.Tensor, + logical_positions: torch.Tensor, + compressor_state_size: int, + query_start_loc: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Map each request to its fixed physical block as a circular token ring.""" + + if compressor_state_size <= 0: + raise ValueError("QSA circular buffer size must be positive") + if block_table.ndim != 2: + raise ValueError("QSA block table must be two-dimensional") + + requests = token_to_req.to(device=block_table.device, dtype=torch.long) + positions = logical_positions.to(device=block_table.device, dtype=torch.long) + if not all(block_table.shape): + slots = torch.full_like(positions, PAD_SLOT_ID) + else: + valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0) + safe_requests = requests.clamp(0, block_table.shape[0] - 1) + physical_blocks = block_table[safe_requests, 0].long() + valid &= physical_blocks >= 0 + slots = physical_blocks * compressor_state_size + positions.remainder( + compressor_state_size + ) + slots = torch.where(valid, slots, PAD_SLOT_ID) + + if query_start_loc is not None: + if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2: + raise ValueError("QSA query starts must contain a terminal offset") + query_start_loc = query_start_loc.to(block_table.device) + num_requests = query_start_loc.shape[0] - 1 + safe_requests = requests.clamp(0, num_requests - 1) + request_ends = query_start_loc.index_select(0, safe_requests + 1) + rows = torch.arange(slots.numel(), device=slots.device) + keep = ( + (requests >= 0) + & (requests < num_requests) + & (rows + compressor_state_size >= request_ends) + ) + slots = torch.where(keep, slots, PAD_SLOT_ID) + + slots = slots.to(torch.int64) + if out is not None: + out.fill_(PAD_SLOT_ID) + out[: slots.numel()].copy_(slots) + return out[: slots.numel()] + return slots + + +def compressed_qsa_slot_mapping( + block_table: torch.Tensor, + token_to_req: torch.Tensor, + logical_positions: torch.Tensor, + storage_block_size: int, + compress_ratio: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Build boundary-only slots for an ``MLAAttentionSpec`` QSA cache.""" + + if storage_block_size <= 0 or compress_ratio <= 0: + raise ValueError("QSA block size and compression ratio must be positive") + compressed_positions = torch.div( + logical_positions.clamp_min(0), compress_ratio, rounding_mode="floor" + ) + slots = _logical_to_physical_qsa_slots( + block_table, + token_to_req, + compressed_positions, + storage_block_size, + ) + valid = (logical_positions >= 0) & ( + (logical_positions + 1).remainder(compress_ratio) == 0 + ) + slots = torch.where(valid, slots, PAD_SLOT_ID).to(torch.int64) + if out is not None: + out.fill_(PAD_SLOT_ID) + out[: slots.numel()].copy_(slots) + return out[: slots.numel()] + return slots + + +@cache +def _metadata_launch_pdl() -> bool: + return current_platform.is_arch_support_pdl() + + +@triton.jit( + do_not_specialize=[ + "num_reqs", + "num_mapped_tokens", + "num_tokens", + "max_num_work", + "num_search_steps", + "work_search_steps", + ] +) +def _build_qsa_metadata_kernel( + query_start_loc_ptr, + seq_lens_ptr, + common_slot_mapping_ptr, + block_table_ptr, + token_to_req_ptr, + logical_positions_ptr, + slot_mapping_ptr, + k_work_metadata_ptr, + block_table_stride_0: tl.constexpr, + block_table_stride_1: tl.constexpr, + num_reqs, + num_mapped_tokens, + num_tokens, + max_num_work, + num_search_steps, + work_search_steps, + storage_block_size: tl.constexpr, + compress_ratio: tl.constexpr, + circular_buffer_size: tl.constexpr, + num_block_table_columns: tl.constexpr, + launch_pdl: tl.constexpr, + TOKEN_BLOCK_SIZE: tl.constexpr, + REQUEST_SCAN_SIZE: tl.constexpr, + WORK_BLOCK_SIZE: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + + pid = tl.program_id(0) + token_idx = pid * TOKEN_BLOCK_SIZE + tl.arange(0, TOKEN_BLOCK_SIZE) + store_mask = token_idx < num_tokens + mapped = token_idx < num_mapped_tokens + search_token_idx = tl.minimum(token_idx, num_mapped_tokens - 1) + request_idx = tl.zeros((TOKEN_BLOCK_SIZE,), tl.int32) + # Find the last query start at or before each token. The dynamic loop avoids + # compiling a kernel variant for every ceil(log2(num_reqs)). + for step in tl.range(0, num_search_steps): + candidate = request_idx + (1 << (num_search_steps - step - 1)) + valid_candidate = candidate < num_reqs + candidate_start = tl.load( + query_start_loc_ptr + candidate, + mask=valid_candidate, + other=num_mapped_tokens + 1, + ) + advance = valid_candidate & (candidate_start <= search_token_idx) + request_idx = tl.where(advance, candidate, request_idx) + query_start = tl.load(query_start_loc_ptr + request_idx, mask=mapped, other=0) + query_end = tl.load(query_start_loc_ptr + request_idx + 1, mask=mapped, other=0) + seq_len = tl.load(seq_lens_ptr + request_idx, mask=mapped, other=0) + logical_position = seq_len - (query_end - query_start) + token_idx - query_start + logical_position = tl.where(mapped, logical_position, -1) + tl.store( + token_to_req_ptr + token_idx, + tl.where(mapped, request_idx, 0), + mask=store_mask, + ) + tl.store( + logical_positions_ptr + token_idx, + logical_position, + mask=store_mask, + ) + + # circular_buffer_size is constexpr, so each builder instance compiles out + # the other QSA cache owner's slot-mapping rule. + if circular_buffer_size > 0: + valid = ( + mapped + & (logical_position >= 0) + & (token_idx + circular_buffer_size >= query_end) + & (num_block_table_columns > 0) + ) + physical_block = tl.load( + block_table_ptr + request_idx * block_table_stride_0, + mask=valid, + other=-1, + ) + valid &= physical_block >= 0 + slot = physical_block * circular_buffer_size + ( + logical_position % circular_buffer_size + ) + elif compress_ratio != 1: + compressed_position = tl.maximum(logical_position, 0) // compress_ratio + logical_block = compressed_position // storage_block_size + valid = ( + mapped + & (logical_position >= 0) + & ((logical_position + 1) % compress_ratio == 0) + & (logical_block < num_block_table_columns) + ) + physical_block = tl.load( + block_table_ptr + + request_idx * block_table_stride_0 + + logical_block * block_table_stride_1, + mask=valid, + other=-1, + ) + valid &= physical_block >= 0 + valid &= ( + tl.load(common_slot_mapping_ptr + token_idx, mask=mapped, other=-1) >= 0 + ) + slot = physical_block * storage_block_size + ( + compressed_position % storage_block_size + ) + if (circular_buffer_size > 0) or (compress_ratio != 1): + tl.store( + slot_mapping_ptr + token_idx, + tl.where(valid, slot, -1), + mask=store_mask, + ) + work_tile_start = pid * WORK_BLOCK_SIZE + has_work_tile = work_tile_start < max_num_work + if k_work_metadata_ptr is not None and has_work_tile: + # Every work CTA builds the request prefix in registers. Recomputing this + # small vector lets CTAs write disjoint work tiles without a grid barrier. + requests = tl.arange(0, REQUEST_SCAN_SIZE) + valid_request = requests < num_reqs + request_query_start = tl.load( + query_start_loc_ptr + requests, mask=valid_request, other=0 + ) + request_query_end = tl.load( + query_start_loc_ptr + requests + 1, mask=valid_request, other=0 + ) + request_seq_len = tl.load(seq_lens_ptr + requests, mask=valid_request, other=0) + request_query_len = request_query_end - request_query_start + chunk_start = request_seq_len - request_query_len + num_groups = request_seq_len // compress_ratio - chunk_start // compress_ratio + # Nonempty requests need one item even without a completed compression + # group because work item zero also commits the current raw-K suffix. + work_counts = tl.where(request_query_len > 0, tl.maximum(num_groups, 1), 0) + work_ends = tl.cumsum(tl.where(valid_request, work_counts, 0), axis=0) + total_work = tl.sum(tl.where(valid_request, work_counts, 0), axis=0) + work_offsets = tl.arange(0, WORK_BLOCK_SIZE) + work = work_tile_start + work_offsets + in_bounds = work < max_num_work + active = in_bounds & (work < total_work) + request = tl.zeros((WORK_BLOCK_SIZE,), dtype=tl.int32) + # Request zero starts at zero for every active work item. Descending + # steps find the last request starting at or before this item. + for step_idx in tl.range(0, work_search_steps): + step = 1 << (work_search_steps - step_idx - 1) + candidate = request + step + valid_candidate = candidate < num_reqs + candidate_start = tl.gather(work_ends, candidate - 1, 0) + advance = active & valid_candidate & (candidate_start <= work) + request = tl.where(advance, candidate, request) + + if launch_pdl: + # Let the dependent grid start launch setup while these CTAs finish + # stores; its gdc_wait still orders access to the completed metadata. + tl.extra.cuda.gdc_launch_dependents() + + owner_work_start = tl.gather(work_ends, tl.maximum(request - 1, 0), 0) + owner_work_start = tl.where(request == 0, 0, owner_work_start) + work_in_request = work - owner_work_start + tl.store( + k_work_metadata_ptr + work * 2, + tl.where(active, request, -1), + mask=in_bounds, + ) + tl.store( + k_work_metadata_ptr + work * 2 + 1, + tl.where(active, work_in_request, -1), + mask=in_bounds, + ) + + else: + if launch_pdl: + # A dependent grid launches only after every CTA has signaled. + tl.extra.cuda.gdc_launch_dependents() + + +def build_qsa_metadata_triton( + common_attn_metadata: CommonAttentionMetadata, + token_to_req_buffer: torch.Tensor, + logical_positions_buffer: torch.Tensor, + slot_mapping_buffer: torch.Tensor, + *, + storage_block_size: int, + compress_ratio: int, + circular_buffer_size: int = 0, + k_work_metadata_buffer: torch.Tensor | None = None, + request_capacity: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build QSA side-cache and optional pre-indexer work metadata.""" + num_tokens = common_attn_metadata.num_actual_tokens + num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) + token_to_req = token_to_req_buffer[:num_tokens] + logical_positions = logical_positions_buffer[:num_tokens] + slot_mapping = slot_mapping_buffer[:num_tokens] + num_reqs = common_attn_metadata.query_start_loc.shape[0] - 1 + assert num_reqs > 0 + + if k_work_metadata_buffer is not None: + if request_capacity is None: + request_capacity = num_reqs + assert request_capacity >= num_reqs + # Pad for tl.arange while keeping the scan width stable across live batches. + request_scan_size = 1 << int(math.ceil(math.log2(request_capacity))) + max_num_work = k_work_metadata_buffer.shape[0] + else: + request_scan_size = 1 + max_num_work = 0 + + if num_tokens == 0 and k_work_metadata_buffer is None: + return token_to_req, logical_positions, slot_mapping + + block_table = common_attn_metadata.block_table_tensor + num_search_steps = int(math.ceil(math.log2(num_reqs))) + work_search_steps = int(math.ceil(math.log2(num_reqs))) + # The same grid covers token tiles and, for the compressed cache, work tiles. + num_token_blocks = cdiv(num_tokens, 128) + num_work_blocks = ( + cdiv(max_num_work, 256) if k_work_metadata_buffer is not None else 0 + ) + _build_qsa_metadata_kernel[(max(num_token_blocks, num_work_blocks, 1),)]( + common_attn_metadata.query_start_loc, + common_attn_metadata.seq_lens, + common_attn_metadata.slot_mapping, + block_table, + token_to_req, + logical_positions, + slot_mapping, + k_work_metadata_buffer, + block_table.stride(0), + block_table.stride(1), + num_reqs, + num_mapped_tokens, + num_tokens, + max_num_work, + num_search_steps, + work_search_steps, + storage_block_size, + compress_ratio, + circular_buffer_size, + block_table.shape[1], + launch_pdl=_metadata_launch_pdl(), + TOKEN_BLOCK_SIZE=128, + REQUEST_SCAN_SIZE=request_scan_size, + WORK_BLOCK_SIZE=256, + num_warps=4, + ) + if circular_buffer_size == 0 and compress_ratio == 1: + slot_mapping = common_attn_metadata.slot_mapping[:num_tokens] + return token_to_req, logical_positions, slot_mapping + + +def _build_qsa_metadata_torch( + common_attn_metadata: CommonAttentionMetadata, + token_to_req_buffer: torch.Tensor, + logical_positions_buffer: torch.Tensor, + slot_mapping_buffer: torch.Tensor, + *, + storage_block_size: int, + compress_ratio: int, + circular_buffer_size: int = 0, + k_work_metadata_buffer: torch.Tensor | None = None, + request_capacity: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del request_capacity + num_tokens = common_attn_metadata.num_actual_tokens + num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) + logical_positions = logical_positions_buffer[:num_tokens] + + token_to_req = common_attn_metadata.token_to_req_indices(token_to_req_buffer)[ + :num_tokens + ] + logical_positions[:num_mapped_tokens].copy_( + _logical_positions( + common_attn_metadata.query_start_loc, + common_attn_metadata.seq_lens, + token_to_req[:num_mapped_tokens], + num_mapped_tokens, + ) + ) + if num_mapped_tokens < num_tokens: + logical_positions[num_mapped_tokens:].fill_(-1) + if circular_buffer_size > 0: + slot_mapping = circular_qsa_slot_mapping( + common_attn_metadata.block_table_tensor, + token_to_req, + logical_positions, + circular_buffer_size, + query_start_loc=common_attn_metadata.query_start_loc, + out=slot_mapping_buffer, + ) + elif compress_ratio == 1: + slot_mapping = common_attn_metadata.slot_mapping[:num_tokens] + else: + slot_mapping = compressed_qsa_slot_mapping( + common_attn_metadata.block_table_tensor, + token_to_req, + logical_positions, + storage_block_size, + compress_ratio, + slot_mapping_buffer, + ) + slot_mapping.masked_fill_( + common_attn_metadata.slot_mapping[:num_tokens] < 0, -1 + ) + if k_work_metadata_buffer is not None: + query_lens = ( + common_attn_metadata.query_start_loc[1:] + - common_attn_metadata.query_start_loc[:-1] + ) + chunk_starts = common_attn_metadata.seq_lens - query_lens + num_work_per_request = ( + common_attn_metadata.seq_lens // compress_ratio + - chunk_starts // compress_ratio + ) + num_work_per_request = torch.where( + query_lens > 0, num_work_per_request.clamp_min(1), 0 + ) + k_start_loc = torch.empty( + query_lens.shape[0] + 1, + dtype=torch.int32, + device=query_lens.device, + ) + k_start_loc[0] = 0 + torch.cumsum(num_work_per_request, 0, out=k_start_loc[1:]) + work = torch.arange( + k_work_metadata_buffer.shape[0], + device=k_work_metadata_buffer.device, + ) + requests = torch.searchsorted(k_start_loc[1:], work, right=True) + active = work < k_start_loc[-1] + work_in_request = ( + work - k_start_loc[requests.clamp_max(query_lens.shape[0] - 1)] + ) + k_work_metadata_buffer[:, 0].copy_( + torch.where(active, requests, -1).to(torch.int32) + ) + k_work_metadata_buffer[:, 1].copy_( + torch.where(active, work_in_request, -1).to(torch.int32) + ) + return token_to_req, logical_positions, slot_mapping + + +# Resolve the fallback outside the per-step metadata hot path. +build_qsa_metadata = ( + build_qsa_metadata_triton if HAS_TRITON else _build_qsa_metadata_torch +) + + +@dataclass +class QSAForwardMetadata(AttentionMetadata): + """Common per-forward metadata for one QSA side cache.""" + + block_table: torch.Tensor + slot_mapping: torch.Tensor + seq_lens: torch.Tensor + query_start_loc: torch.Tensor + token_to_req: torch.Tensor + logical_positions: torch.Tensor + k_work_metadata: torch.Tensor + num_actual_tokens: int + storage_block_size: int + compress_ratio: int + + +class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]): + """Build QSA metadata from vLLM's cache-group-specific common metadata.""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.is_circular_buffer = isinstance(kv_cache_spec, CircularBufferSpec) + if isinstance(kv_cache_spec, MLAAttentionSpec): + self.compress_ratio = kv_cache_spec.compress_ratio + else: + self.compress_ratio = 1 + self.storage_block_size = kv_cache_spec.storage_block_size + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.token_to_req_buffer = torch.empty( + max_tokens, dtype=torch.int32, device=device + ) + self.slot_mapping_buffer = torch.empty( + max_tokens, dtype=torch.int64, device=device + ) + self.logical_positions_buffer = torch.empty( + max_tokens, dtype=torch.int64, device=device + ) + max_requests = vllm_config.scheduler_config.max_num_seqs + self.request_capacity = max_requests + if not self.is_circular_buffer and self.compress_ratio != 1: + max_k_work = ( + max_tokens + (self.compress_ratio - 1) * max_requests + ) // self.compress_ratio + self.k_work_metadata_buffer = torch.empty( + max_k_work, 2, dtype=torch.int32, device=device + ) + else: + self.k_work_metadata_buffer = torch.empty( + 0, 2, dtype=torch.int32, device=device + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> QSAForwardMetadata: + del common_prefix_len, fast_build + num_tokens = common_attn_metadata.num_actual_tokens + build_k_work = not self.is_circular_buffer and self.compress_ratio != 1 + k_work_metadata = self.k_work_metadata_buffer + request_capacity = None + if build_k_work: + num_requests = common_attn_metadata.query_start_loc.shape[0] - 1 + request_capacity = self.request_capacity + max_num_work = ( + num_tokens + (self.compress_ratio - 1) * num_requests + ) // self.compress_ratio + k_work_metadata = self.k_work_metadata_buffer[:max_num_work] + token_to_req, logical_positions, slot_mapping = build_qsa_metadata( + common_attn_metadata, + self.token_to_req_buffer, + self.logical_positions_buffer, + self.slot_mapping_buffer, + storage_block_size=self.storage_block_size, + compress_ratio=self.compress_ratio, + circular_buffer_size=( + self.kv_cache_spec.block_size if self.is_circular_buffer else 0 + ), + k_work_metadata_buffer=k_work_metadata if build_k_work else None, + request_capacity=request_capacity, + ) + return QSAForwardMetadata( + block_table=common_attn_metadata.block_table_tensor, + slot_mapping=slot_mapping, + seq_lens=common_attn_metadata.seq_lens, + query_start_loc=common_attn_metadata.query_start_loc, + token_to_req=token_to_req, + logical_positions=logical_positions, + k_work_metadata=k_work_metadata, + num_actual_tokens=num_tokens, + storage_block_size=self.storage_block_size, + compress_ratio=self.compress_ratio, + ) + + +class QSAStateBackend(AttentionBackend): + """Key-only dummy backend for out-of-band QSA side-cache operations.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [ + torch.float16, + torch.bfloat16, + ] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "float16", + "bfloat16", + ] + + @staticmethod + def get_name() -> str: + return "QWEN4_EXP_EXP_QSA_STATE" + + @staticmethod + def get_impl_cls(): + raise NotImplementedError( + "QSA state caches run out-of-band and have no attention impl" + ) + + @staticmethod + def get_builder_cls() -> type[QSAMetadataBuilder]: + return QSAMetadataBuilder + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + del cache_dtype_str + if num_kv_heads != 1: + raise ValueError("QSA side caches require exactly one KV head") + return (num_blocks, block_size, num_kv_heads, head_size) + + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + return True + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + if include_num_layers_dimension: + return (0, 1, 2, 3, 4) + return (0, 1, 2, 3) + + +class _QSAStateCache(nn.Module, AttentionLayerBase): + supports_dcp = False + + def __init__( + self, + *, + head_size: int, + dtype: torch.dtype, + cache_config: CacheConfig, + prefix: str, + vllm_config: VllmConfig, + compress_ratio: int = 1, + ) -> None: + super().__init__() + if head_size <= 0: + raise ValueError("QSA cache head size must be positive") + if compress_ratio <= 0: + raise ValueError("QSA compression ratio must be positive") + if cache_config.block_size % compress_ratio: + raise ValueError( + "QSA cache block size must be divisible by the compression ratio" + ) + self.head_size = head_size + self.dtype = dtype + self.cache_config = cache_config + self.prefix = prefix + self.compress_ratio = compress_ratio + self.kv_cache = torch.tensor([]) + + static_context = vllm_config.compilation_config.static_forward_context + if prefix in static_context: + raise ValueError(f"Duplicate layer name: {prefix}") + static_context[prefix] = self + + def forward(self) -> None: ... + + def get_attn_backend(self) -> type[AttentionBackend]: + return QSAStateBackend + + +class QSAKeyStateCache(_QSAStateCache): + """Raw 16-bit key, optionally followed by exact int64 MRoPE positions.""" + + _BF16_PER_INT64 = 4 + _NUM_ROPE_AXES = 3 + + def __init__(self, *, cache_rope_positions: bool = False, **kwargs) -> None: + key_head_size = int(kwargs.pop("head_size")) + self.key_head_size = key_head_size + self.cache_rope_positions = bool(cache_rope_positions) + self.rope_position_offset = ( + (key_head_size + self._BF16_PER_INT64 - 1) // self._BF16_PER_INT64 + ) * self._BF16_PER_INT64 + storage_head_size = key_head_size + if self.cache_rope_positions: + storage_head_size = self.rope_position_offset + ( + self._NUM_ROPE_AXES * self._BF16_PER_INT64 + ) + super().__init__(head_size=storage_head_size, **kwargs) + + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + if kv_cache.ndim != 4 or kv_cache.shape[2] != 1: + raise ValueError("QSA raw cache must be [blocks, block_size, 1, width]") + if kv_cache.dtype != self.dtype or kv_cache.shape[3] != self.head_size: + raise ValueError("QSA raw cache does not match its packed cache spec") + super().bind_kv_cache(kv_cache) + self.key_cache = kv_cache[..., : self.key_head_size] + if self.cache_rope_positions: + position_tail = kv_cache[..., self.rope_position_offset :] + self.rope_position_cache = position_tail.view(torch.int64) + else: + self.rope_position_cache = None + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + # Hold the open group's committed keys plus every row a speculative + # step stores before acceptance is known, rounded up to whole groups so + # the ring divides the attention block size (it joins the LCM that sets + # the scheduler block size). Anything narrower lets a rejected draft row + # overwrite a committed key the next step needs to close the group. + span = self.compress_ratio + vllm_config.num_speculative_tokens + capacity = self.compress_ratio * cdiv(span, self.compress_ratio) + assert self.cache_config.block_size % capacity == 0, ( + f"QSA ring capacity {capacity} must divide the attention block " + f"size {self.cache_config.block_size}" + ) + return CircularBufferSpec( + block_size=capacity, + num_kv_heads=1, + head_size=self.head_size, + head_size_v=0, + dtype=self.dtype, + ) + + +class QSACompressedKeyCache(_QSAStateCache): + """Normalized, group-first-RoPE 16-bit key per complete group.""" + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + del vllm_config + return MLAAttentionSpec( + block_size=self.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_size, + dtype=self.dtype, + compress_ratio=self.compress_ratio, + ) + + +__all__ = [ + "QSACompressedKeyCache", + "QSAForwardMetadata", + "QSAKeyStateCache", + "QSAMetadataBuilder", + "QSAStateBackend", + "canonical_qsa_rope_positions", + "circular_qsa_slot_mapping", + "compressed_qsa_slot_mapping", +] diff --git a/vllm/models/qwen4_exp/config.py b/vllm/models/qwen4_exp/config.py new file mode 100644 index 0000000000..a8eedab83f --- /dev/null +++ b/vllm/models/qwen4_exp/config.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp model configuration.""" + +from typing import Any, ClassVar, cast + +from transformers import PretrainedConfig +from transformers.models.qwen3_vl.configuration_qwen3_vl import ( + Qwen3VLVisionConfig, +) + +from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig + +_QSA_CONFIG_FIELDS = ( + "indexer_n_heads", + "indexer_kv_heads", + "indexer_head_dim", + "indexer_budget", + "indexer_compress_ratio", +) + + +class Qwen4ExpVisionConfig(Qwen3VLVisionConfig): + model_type = "qwen4_exp" + base_config_key = "vision_config" + + +class Qwen4ExpTextConfig(Qwen3NextConfig): + model_type = "qwen4_exp_text" + base_config_key = "text_config" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + hc_count: int = 4, + hc_lowrank: int = 320, + ple_layer_ids: list[int] | None = None, + ple_embed_dim: int | None = None, + ple_conv_kernel_size: int = 4, + ngram_size: int = 3, + heads_per_ngram: int = 8, + ngram_vocab_size_base: int = 20_000_000, + make_ngram_vocab_size_divisible_by: int = 128, + output_gate_type: str = "sigmoid", + rope_parameters: dict[str, Any] | None = None, + layer_types: list[str] | None = None, + **kwargs: Any, + ) -> None: + if hc_count <= 1: + raise ValueError(f"Qwen4Exp requires hc_count > 1, got {hc_count}.") + + if rope_parameters is not None: + if kwargs.get("rope_scaling") is None: + kwargs["rope_scaling"] = rope_parameters + if kwargs.get("rope_theta") is None and "rope_theta" in rope_parameters: + kwargs["rope_theta"] = rope_parameters["rope_theta"] + if ( + kwargs.get("partial_rotary_factor") is None + and "partial_rotary_factor" in rope_parameters + ): + kwargs["partial_rotary_factor"] = rope_parameters[ + "partial_rotary_factor" + ] + + rope_scaling = kwargs.get("rope_scaling") + rope_theta = kwargs.get("rope_theta", 10_000.0) + super().__init__(layer_types=layer_types, **kwargs) + + normalized_rope_parameters = self.rope_parameters + self.rope_scaling = ( + rope_scaling or rope_parameters or normalized_rope_parameters + ) + self.rope_parameters = rope_parameters or normalized_rope_parameters + self.rope_theta = rope_theta + + self.hc_count = hc_count + self.hc_lowrank = hc_lowrank + self.ple_layer_ids = ple_layer_ids or [] + self.ple_embed_dim = ( + self.hidden_size if ple_embed_dim is None else ple_embed_dim + ) + self.ple_conv_kernel_size = ple_conv_kernel_size + self.ngram_size = ngram_size + self.heads_per_ngram = heads_per_ngram + self.ngram_vocab_size_base = ngram_vocab_size_base + self.make_ngram_vocab_size_divisible_by = make_ngram_vocab_size_divisible_by + self.output_gate_type = output_gate_type + + self._validate_ple_config() + self._validate_ple_layer_ids() + self._validate_qsa_config() + + def _validate_ple_config(self) -> None: + if self.hc_lowrank <= 0: + raise ValueError(f"hc_lowrank must be positive, got {self.hc_lowrank}") + if self.ngram_size < 2: + raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}") + if self.heads_per_ngram <= 0: + raise ValueError( + f"heads_per_ngram must be positive, got {self.heads_per_ngram}" + ) + if self.ple_embed_dim <= 0: + raise ValueError( + f"ple_embed_dim must be positive, got {self.ple_embed_dim}" + ) + ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram + if self.ple_embed_dim % ngram_heads: + raise ValueError( + "ple_embed_dim must be divisible by total ngram heads: " + f"{self.ple_embed_dim} % {ngram_heads} != 0" + ) + if self.ple_conv_kernel_size <= 0: + raise ValueError( + "ple_conv_kernel_size must be positive, got " + f"{self.ple_conv_kernel_size}" + ) + if self.ngram_vocab_size_base <= 0: + raise ValueError("ngram_vocab_size_base must be positive") + if self.make_ngram_vocab_size_divisible_by <= 0: + raise ValueError("make_ngram_vocab_size_divisible_by must be positive") + + def _validate_ple_layer_ids(self) -> None: + invalid = [ + layer_id + for layer_id in self.ple_layer_ids + if not 1 <= int(layer_id) <= self.num_hidden_layers + ] + if invalid: + raise ValueError( + "ple_layer_ids are 1-based and must refer to an existing layer; " + f"got {invalid} for {self.num_hidden_layers} layers" + ) + + def _validate_qsa_config(self) -> None: + configured = {name: getattr(self, name, None) for name in _QSA_CONFIG_FIELDS} + if all(value is None for value in configured.values()): + return + + missing = [name for name, value in configured.items() if value is None] + if missing: + raise ValueError(f"QSA config is missing required fields: {missing}") + + values = {name: int(cast(int, value)) for name, value in configured.items()} + if any(value <= 0 for value in values.values()): + raise ValueError(f"QSA config values must be positive: {values}") + if values["indexer_kv_heads"] != 1: + raise ValueError("the QSA MQA operators require indexer_kv_heads=1") + if values["indexer_budget"] % values["indexer_compress_ratio"] != 0: + raise ValueError( + "indexer_budget must be divisible by indexer_compress_ratio" + ) + block_topk = values["indexer_budget"] // values["indexer_compress_ratio"] + if block_topk not in (512, 2048): + raise ValueError( + "QSA requires indexer_budget / indexer_compress_ratio " + f"to be 512 or 2048, got {block_topk}" + ) + rotary_dim = int(self.head_dim * self.partial_rotary_factor) + if rotary_dim > values["indexer_head_dim"]: + raise ValueError( + "QSA indexer_head_dim must cover the attention rotary " + f"dimension, got {values['indexer_head_dim']} < {rotary_dim}" + ) + + @property + def layers_block_type(self) -> list[str]: + return [ + "attention" if layer_type == "full_attention" else layer_type + for layer_type in self.layer_types + ] + + @property + def short_conv_layer_ids(self) -> list[int]: + if not self.ple_layer_ids: + return [] + return sorted({int(layer_id) - 1 for layer_id in self.ple_layer_ids}) + + @property + def short_conv_state_shape(self) -> tuple[int, int] | None: + if not self.short_conv_layer_ids: + return None + ple_state_len = (self.ple_conv_kernel_size - 1) * self.ngram_size + ple_channels = self.hidden_size * self.hc_count + return ple_channels, ple_state_len + + @property + def ngram_context_len(self) -> int: + if not self.ple_layer_ids: + return 0 + return max(int(self.ngram_size) - 1, 0) + + +class Qwen4ExpConfig(PretrainedConfig): + model_type = "qwen4_exp" + sub_configs: ClassVar[dict[str, type[PretrainedConfig]]] = { + "vision_config": Qwen4ExpVisionConfig, + "text_config": Qwen4ExpTextConfig, + } + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config: Qwen4ExpTextConfig | dict[str, Any] | None = None, + vision_config: Qwen4ExpVisionConfig | dict[str, Any] | None = None, + image_token_id: int = 248056, + video_token_id: int = 248057, + vision_start_token_id: int = 248053, + vision_end_token_id: int = 248054, + tie_word_embeddings: bool = False, + rope_parameters: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + if text_config is not None: + kwargs.pop("split_ngram_parts", None) + + text_kwargs = ( + dict(kwargs) + if text_config is None + and "hidden_size" in kwargs + and "num_hidden_layers" in kwargs + else {} + ) + + if isinstance(vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**vision_config) + elif vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + else: + self.vision_config = vision_config + + if isinstance(text_config, dict): + self.text_config = self.sub_configs["text_config"](**text_config) + elif text_config is None: + self.text_config = self.sub_configs["text_config"](**text_kwargs) + else: + self.text_config = text_config + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id + self.vision_end_token_id = vision_end_token_id + self.rope_parameters = rope_parameters or getattr( + self.text_config, "rope_parameters", {} + ) + super().__init__(**kwargs, tie_word_embeddings=tie_word_embeddings) + + +__all__ = [ + "Qwen4ExpConfig", + "Qwen4ExpTextConfig", + "Qwen4ExpVisionConfig", +] diff --git a/vllm/models/qwen4_exp/nvidia/__init__.py b/vllm/models/qwen4_exp/nvidia/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/qwen4_exp/nvidia/hyperconnection.py b/vllm/models/qwen4_exp/nvidia/hyperconnection.py new file mode 100644 index 0000000000..dbe637d80d --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/hyperconnection.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HyperConnection (Gated Residual) utilities — NVIDIA model variant. + +Implements the HyperConnection residual scheme proposed in +"HyperConnections" (https://arxiv.org/abs/2409.19606). This NVIDIA variant +delays each HC combine to the following HC mix boundary. HC glue kernels, +including fused combine+RMSNorm, live in ``ops/hc.py``; projections remain +standard vLLM Linear modules. + +Hidden states between layers have shape ``[..., HC*HS]`` with HS inner +(HC outer, HS inner — checkpoint-native layout). + +Typical usage inside a transformer decoder layer:: + + self.attn_hc = GatedResidual(hc_config) + + hidden_states, block_input, injection = self.attn_hc.mix(hidden_states) + attention_output = attention(block_input) + hidden_states, block_input, injection = self.mlp_hc.combine_and_mix( + hidden_states, attention_output, injection + ) +""" + +import torch +from torch import nn + +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + ReplicatedLinear, +) +from vllm.model_executor.layers.quantization.sm70_online_qpn8 import ( + maybe_apply_fused_hc, +) +from vllm.model_executor.models.utils import maybe_prefix + +from ..common.hyperconnection import ( + GroupedGemmaRMSNorm, + HyperConnectionConfig, +) +from .ops.hc import ( + grouped_gemma_rmsnorm, + hc_combine, + hc_combine_norm, + hc_gate_mix, + hc_silu, +) + + +# --------------------------------------------------------------------------- +# Gated-residual variant +# --------------------------------------------------------------------------- +class GatedResidual(nn.Module): + """Gated HyperConnection with learnable low-rank mixing and injection. + + ``combine_and_mix()`` runs the pre pipeline (grouped GemmaRMSNorm -> merged + low-rank down+inject GEMM -> silu -> up GEMM -> sigmoid -> gated mean + over the HC streams). When passed a pending block output and an injection, + it fuses their residual combine with the RMSNorm. Final mixers use + ``use_combine=False`` and do not produce a new injection. + + Weights: the norm owns the grouped GemmaRMSNorm affine; the projections + are vLLM Linear modules (merged replicated linear for down+inject), so + GEMM dispatch (e.g. the low-latency skinny GEMM) applies through the + standard quant_method mechanism. + """ + + def __init__( + self, + config: HyperConnectionConfig, + use_combine: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.lora_rank = config.hc_lowrank + self.hc_count = config.hc_count + self.hidden_size = config.hidden_size + self.use_combine = use_combine + + norm_size = ( + self.hyper_hidden_size if config.hc_per_branch_norm else config.hidden_size + ) + group_size = config.hidden_size if config.hc_per_branch_norm else None + # Normalize each H-sized HC stream independently while retaining a + # separate affine weight for every element of the HC*H layout. + self.hc_norm = GroupedGemmaRMSNorm( + norm_size, + eps=config.rms_norm_eps, + group_size=group_size, + dtype=config.params_dtype, + ) + + # -- vLLM Linear weights -------------------------------------------- + # The merged skinny-GEMM shape is physically padded to 16 rows to ensure + # good alignment and performant implementation chosen by CuBLAS heuristics. + self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0 + if use_combine: + self.input_mix_weight_down_block_inject = MergedColumnParallelLinear( + self.hyper_hidden_size, + [self.lora_rank, self.hc_count] + + ([self.pad_size] if self.pad_size else []), + bias=False, + params_dtype=config.params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "input_mix_weight_down_block_inject"), + return_bias=False, + disable_tp=True, + ) + else: + self.input_mix_weight_down = ReplicatedLinear( + self.hyper_hidden_size, + self.lora_rank, + bias=False, + params_dtype=config.params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "input_mix_weight_down"), + return_bias=False, + ) + self.input_mix_weight_up = ReplicatedLinear( + self.lora_rank, + self.hyper_hidden_size, + bias=False, + params_dtype=config.params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "input_mix_weight_up"), + return_bias=False, + ) + + def _project(self, xn: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: + if self.use_combine: + fused = maybe_apply_fused_hc( + self.input_mix_weight_down_block_inject, + self.input_mix_weight_up, + xn, + ) + if fused is not None: + return fused + + split_sizes = [self.lora_rank, self.hc_count, self.pad_size] + down_and_injection = self.input_mix_weight_down_block_inject(xn) + lora, injection, _ = down_and_injection.split(split_sizes, dim=-1) + else: + lora = self.input_mix_weight_down(xn) + injection = None + + lora = hc_silu(lora, self.hc_count) + gate = self.input_mix_weight_up(lora) + return hc_gate_mix(xn, gate, self.hc_count), injection + + def mix( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + xn = grouped_gemma_rmsnorm( + hidden_states, + self.hc_norm.weight, + self.config.rms_norm_eps, + self.hc_count, + ) + + block_input, injection = self._project(xn) + + return hidden_states, block_input, injection + + def combine_and_mix( + self, + hidden_states: torch.Tensor, + prev_block_output: torch.Tensor, + prev_injection: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Consume a pending combine, then prepare the next block input. + + ``hidden_states`` is the multi-stream state from before the pending + block's mix. Its combine with ``block_output`` is fused with this + module's input RMSNorm. + """ + hidden_states, xn = hc_combine_norm( + hidden_states, + prev_block_output, + prev_injection, + self.hc_norm.weight, + self.config.rms_norm_eps, + self.hc_count, + ) + + block_input, injection = self._project(xn) + + return hidden_states, block_input, injection + + def combine( + self, + hidden_states: torch.Tensor, + block_output: torch.Tensor, + injection: torch.Tensor, + ) -> torch.Tensor: + return hc_combine(hidden_states, block_output, injection, self.hc_count) + + @property + def hyper_hidden_size(self) -> int: + return self.hc_count * self.hidden_size + + +__all__ = [ + "GatedResidual", + "GroupedGemmaRMSNorm", + "HyperConnectionConfig", +] diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py new file mode 100644 index 0000000000..e704624146 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp weight-free QSA indexer.""" + +from __future__ import annotations + +from typing import cast + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope +from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, +) + +from ..common.qsa_cache import ( + QSACompressedKeyCache, + QSAForwardMetadata, + QSAKeyStateCache, + canonical_qsa_rope_positions, +) +from .ops.qsa_pre_indexer import qsa_pre_indexer + + +def apply_qsa_rope( + rotary_emb: nn.Module, + positions: torch.Tensor, + tensor: torch.Tensor, +) -> torch.Tensor: + """Apply the main attention's exact 1D/MRoPE composition to QSA heads.""" + + num_tokens, _, head_dim = tensor.shape + rotary_dim = rotary_emb.rotary_dim + cache = rotary_emb._match_cos_sin_cache_dtype(tensor) # noqa: SLF001 + cos_sin = cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + if positions.ndim == 2: + shape = tensor.shape + tensor, _ = triton_mrope( + tensor.reshape(num_tokens, -1), + tensor.new_empty((num_tokens, head_dim)), + cos, + sin, + rotary_emb.mrope_section, + head_dim, + rotary_dim, + rotary_emb.mrope_interleaved, + ) + return tensor.reshape(shape) + + rotated = rotary_emb.apply_rotary_emb.forward_cuda( + tensor[..., :rotary_dim], + cos, + sin, + ) + return torch.cat((rotated, tensor[..., rotary_dim:]), dim=-1) + + +def _supports_fused_pre_indexer( + rotary_emb: nn.Module, + head_dim: int, + num_kv_heads: int, + compress_ratio: int, +) -> bool: + rotary_dim = int(rotary_emb.rotary_dim) + mrope_section = getattr(rotary_emb, "mrope_section", None) + return ( + bool(getattr(rotary_emb, "is_neox_style", False)) + and ( + not mrope_section + or ( + len(mrope_section) == 3 + and sum(mrope_section) == rotary_dim // 2 + and bool(getattr(rotary_emb, "mrope_interleaved", False)) + ) + ) + and head_dim == 128 + and rotary_dim == 64 + and num_kv_heads == 1 + and compress_ratio > 1 + and compress_ratio & (compress_ratio - 1) == 0 + ) + + +class QSAIndexer(nn.Module): + """Replicated Q/K projection plus paged, weight-free QSA selection. + + ``prefix`` must be the checkpoint's indexer prefix, normally + ``model.layers.N.self_attn.indexer``. Consequently the trainable names are + ``index_qk_proj``, ``q_layernorm`` and ``k_layernorm`` under that prefix. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + config: Qwen4ExpTextConfig, + layer_id: int, + rotary_emb: nn.Module, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + if vllm_config.cache_config is None: + raise ValueError("QSA requires a paged KV cache") + if vllm_config.model_config.dtype not in (torch.float16, torch.bfloat16): + raise NotImplementedError("Qwen4Exp QSA requires FP16 or BF16") + + self.layer_id = int(layer_id) + self.index_n_heads = int(config.indexer_n_heads) + self.index_kv_heads = int(config.indexer_kv_heads) + self.index_head_dim = int(config.indexer_head_dim) + self.token_topk = int(config.indexer_budget) + self.compress_ratio = int(config.indexer_compress_ratio) + self.rotary_emb = rotary_emb + self.use_fused_pre_indexer = _supports_fused_pre_indexer( + rotary_emb, + self.index_head_dim, + self.index_kv_heads, + self.compress_ratio, + ) + self.prefix = prefix + # MTP step 0 selects the target-aligned rows; later steps reuse them + # while continuing to update the QSA side cache. + self.skip_topk = False + + self.index_qk_proj = ReplicatedLinear( + int(config.hidden_size), + (self.index_n_heads + self.index_kv_heads) * self.index_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.index_qk_proj" if prefix else "index_qk_proj", + ) + self.q_layernorm = GemmaRMSNorm( + self.index_head_dim, + eps=float(getattr(config, "rms_norm_eps", 1e-6)), + ) + self.k_layernorm = GemmaRMSNorm( + self.index_head_dim, + eps=float(getattr(config, "rms_norm_eps", 1e-6)), + ) + + cache_config = vllm_config.cache_config + cache_prefix = f"{prefix}." if prefix else "" + self.raw_key_cache = QSAKeyStateCache( + head_size=self.index_head_dim, + dtype=vllm_config.model_config.dtype, + cache_rope_positions=vllm_config.model_config.uses_mrope, + prefix=f"{cache_prefix}raw_key_cache", + cache_config=cache_config, + compress_ratio=self.compress_ratio, + vllm_config=vllm_config, + ) + self.compressed_key_cache = QSACompressedKeyCache( + head_size=self.index_head_dim, + dtype=vllm_config.model_config.dtype, + compress_ratio=self.compress_ratio, + prefix=f"{cache_prefix}compressed_key_cache", + cache_config=cache_config, + vllm_config=vllm_config, + ) + + @property + def output_width(self) -> int: + return self.token_topk + self.compress_ratio - 1 + + def _metadata( + self, + ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None: + metadata = get_forward_context().attn_metadata + if isinstance(metadata, list): + metadata = metadata[0] + if not isinstance(metadata, dict): + return None + raw = cast(QSAForwardMetadata, metadata[self.raw_key_cache.prefix]) + compressed = cast( + QSAForwardMetadata, metadata[self.compressed_key_cache.prefix] + ) + if raw.num_actual_tokens != compressed.num_actual_tokens: + raise RuntimeError("QSA side-cache metadata token counts disagree") + if not raw.logical_positions.is_cuda and ( + not torch.equal(raw.logical_positions, compressed.logical_positions) + ): + raise RuntimeError("QSA side-cache metadata positions disagree") + return raw, compressed + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Return fixed-width request-relative token indices padded with ``-1``.""" + + metadata = self._metadata() + if metadata is None: + # Preserve step-0 indices when later MTP steps reuse the buffer. + if self.skip_topk and out is not None: + return out + result = torch.full( + (hidden_states.shape[0], self.output_width), + -1, + dtype=torch.int32, + device=hidden_states.device, + ) + if out is not None: + out.copy_(result) + return out + return result + + from .ops.qsa import ( + qsa_compress_groups_with_ratio, + qsa_select_paged_tokens, + qsa_store_cache_rows, + ) + + raw_metadata, compressed_metadata = metadata + num_tokens = raw_metadata.num_actual_tokens + hidden_states = hidden_states[:num_tokens] + positions = positions[..., :num_tokens] + + # Q/K projection + projected_qk, _ = self.index_qk_proj(hidden_states) + projected_q, raw_keys = projected_qk.split( + ( + self.index_n_heads * self.index_head_dim, + self.index_kv_heads * self.index_head_dim, + ), + dim=-1, + ) + raw_key_state_cache = self.raw_key_cache + compressed_key_cache = self.compressed_key_cache.kv_cache + + if self.use_fused_pre_indexer: + q = projected_q.new_empty( + num_tokens, + self.index_n_heads, + self.index_head_dim, + ) + qsa_pre_indexer( + projected_q, + raw_keys, + positions, + self.rotary_emb.cos_sin_cache, + self.q_layernorm.weight, + self.k_layernorm.weight, + self.q_layernorm.variance_epsilon, + q, + raw_key_state_cache.kv_cache, + raw_metadata.slot_mapping, + raw_metadata.block_table, + raw_metadata.query_start_loc, + raw_metadata.logical_positions, + compressed_key_cache, + compressed_metadata.slot_mapping, + compressed_metadata.k_work_metadata, + compress_ratio=self.compress_ratio, + mrope_section=getattr(self.rotary_emb, "mrope_section", None), + rope_pos_offset=( + raw_key_state_cache.rope_position_offset + if raw_key_state_cache.rope_position_cache is not None + else None + ), + ) + else: + # Unfused reference path + from flashinfer.norm import gemma_rmsnorm + + q = projected_q.reshape(-1, self.index_n_heads, self.index_head_dim) + q = gemma_rmsnorm( + q.reshape(-1, self.index_head_dim), + self.q_layernorm.weight, + self.q_layernorm.variance_epsilon, + ).reshape_as(q) + q = apply_qsa_rope(self.rotary_emb, positions, q) + + raw_key_cache = raw_key_state_cache.key_cache + rope_position_cache = raw_key_state_cache.rope_position_cache + if rope_position_cache is None: + position_rows = raw_metadata.logical_positions.view(-1, 1, 1).expand( + -1, 1, 3 + ) + else: + position_rows = canonical_qsa_rope_positions(positions).to( + device=raw_key_cache.device + ) + pooled, first_positions = qsa_compress_groups_with_ratio( + raw_keys.reshape(-1, 1, self.index_head_dim), + position_rows, + raw_key_cache, + raw_metadata.block_table, + raw_metadata.token_to_req, + raw_metadata.query_start_loc, + raw_metadata.logical_positions, + compressed_metadata.slot_mapping, + self.compress_ratio, + rope_position_cache, + ) + compressed_keys = gemma_rmsnorm( + pooled.reshape(-1, self.index_head_dim), + self.k_layernorm.weight, + self.k_layernorm.variance_epsilon, + ).reshape(-1, 1, self.index_head_dim) + if getattr(self.rotary_emb, "mrope_section", None): + first_positions = first_positions.transpose(0, 1) + else: + first_positions = first_positions[:, 0] + compressed_keys = apply_qsa_rope( + self.rotary_emb, + first_positions, + compressed_keys, + ) + qsa_store_cache_rows( + compressed_key_cache, + compressed_metadata.slot_mapping, + compressed_keys, + ) + qsa_store_cache_rows( + raw_key_cache, + raw_metadata.slot_mapping, + raw_keys, + ) + if rope_position_cache is not None: + qsa_store_cache_rows( + rope_position_cache, + raw_metadata.slot_mapping, + position_rows, + ) + + if self.skip_topk: + if out is None: + raise RuntimeError("QSA top-k reuse requires an output buffer") + return out + + # Score compressed keys, select blocks, then expand them to token indices. + return qsa_select_paged_tokens( + q, + compressed_key_cache, + compressed_metadata.block_table, + compressed_metadata.token_to_req, + compressed_metadata.logical_positions, + compressed_metadata.seq_lens, + self.token_topk, + self.compress_ratio, + out, + ) + + +__all__ = ["QSAIndexer", "apply_qsa_rope"] diff --git a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py new file mode 100644 index 0000000000..43b7a76bc5 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp decode GEMM selection on Blackwell. + +Dispatch follows Kimi-K3 and uses the local ``(N, K)`` shape and token count. +Plans contain measured CUDA graph capture sizes; other token counts use the +standard linear implementation. +""" + +import torch +from torch import nn + +import vllm.envs as envs +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + SkinnyGemmConfig, + shape_dynamic_skinny_gemm, +) +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + UnquantizedEmbeddingMethod, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import direct_register_custom_op + +QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = { + # GDN fused QKVZ projection, TP=4. + (4096, 2560): { + 1: SkinnyGemmConfig(1, 64, 4, k_unroll=4), + 2: SkinnyGemmConfig(2, 64, 4, k_unroll=4), + }, + # GDN and QSA output projections, TP=4. + (2560, 1536): { + 1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4), + 2: SkinnyGemmConfig(2, 128, 2, k_unroll=2, vector_width=4), + 4: SkinnyGemmConfig(4, 64, 2, k_unroll=2), + }, + # GDN fused B/A projection, TP=4. + (24, 2560): { + 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4), + 2: SkinnyGemmConfig(2, 128, 2, k_unroll=4, vector_width=4), + 4: SkinnyGemmConfig(4, 128, 2, k_unroll=4, vector_width=4), + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4), + 16: SkinnyGemmConfig(16, 128, 1, k_unroll=4, vector_width=4), + }, + # QSA fused QKV/gate projection, TP=4. + (3584, 2560): { + 1: SkinnyGemmConfig(1, 128, 4, k_unroll=4, vector_width=4), + 2: SkinnyGemmConfig(2, 64, 2, k_unroll=2), + 4: SkinnyGemmConfig(4, 64, 2, k_unroll=2), + }, + # QSA indexer Q/K projection, replicated in a TP=4 deployment. + (640, 2560): { + 1: SkinnyGemmConfig(1, 128, 1, k_unroll=4, vector_width=4), + 2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4), + 4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4), + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4), + }, + # Shared-expert fused gate/up projection, TP=4. + (320, 2560): { + 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4), + 2: SkinnyGemmConfig(2, 128, 2, k_unroll=4, vector_width=4), + 4: SkinnyGemmConfig(4, 128, 2, k_unroll=4, vector_width=4), + 8: SkinnyGemmConfig(8, 64, 1, k_unroll=4), + 16: SkinnyGemmConfig(16, 128, 2, k_unroll=4, vector_width=4), + }, + # LM head, TP=4. + (62080, 2560): { + 1: SkinnyGemmConfig(1, 64, 4, k_unroll=2), + 2: SkinnyGemmConfig(2, 32, 4, k_unroll=2), + }, + # HC merged down/injection projection, replicated in a TP=4 deployment. + (336, 10240): { + 1: SkinnyGemmConfig(1, 128, 1, static_k=10240), + 2: SkinnyGemmConfig(2, 128, 1, static_k=10240), + 4: SkinnyGemmConfig(4, 128, 2, static_k=10240), + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4), + }, +} + + +def _is_sm103() -> bool: + return current_platform.is_device_capability((10, 3)) + + +def _is_packed_row_major(tensor: torch.Tensor) -> bool: + return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1) + + +def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + not envs.VLLM_BATCH_INVARIANT + and _is_packed_row_major(x) + and _is_packed_row_major(weight) + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.is_cuda + and weight.is_cuda + and x.device == weight.device + and x.shape[1] == weight.shape[1] + ) + + +class _Qwen4ExpLowLatencyApply: + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if bias is None and not envs.VLLM_BATCH_INVARIANT: + return torch.ops.vllm.qwen4_exp_low_latency_gemm(x, layer.weight) + return super().apply(layer, x, bias) # type: ignore[misc] + + +class Qwen4ExpLowLatencyLinearMethod(_Qwen4ExpLowLatencyApply, UnquantizedLinearMethod): + pass + + +class Qwen4ExpLowLatencyEmbeddingMethod( + _Qwen4ExpLowLatencyApply, UnquantizedEmbeddingMethod +): + pass + + +def _qwen4_exp_low_latency_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1])) + config = None if plan is None else plan.get(x.shape[0]) + if ( + config is not None + and _runtime_ok(x, weight) + and shape_dynamic_skinny_gemm.is_available() + ): + return shape_dynamic_skinny_gemm(x, weight, config) + return torch.nn.functional.linear(x, weight) + + +def _qwen4_exp_low_latency_gemm_fake( + x: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: + return x.new_empty((*x.shape[:-1], weight.shape[0])) + + +direct_register_custom_op( + op_name="qwen4_exp_low_latency_gemm", + op_func=_qwen4_exp_low_latency_gemm, + fake_impl=_qwen4_exp_low_latency_gemm_fake, +) + + +def enable_qwen4_exp_low_latency_gemm( + module: nn.Module, + dtype: torch.dtype, +) -> None: + if dtype != torch.bfloat16 or not _is_sm103(): + return + if not shape_dynamic_skinny_gemm.is_available(): + return + + warmup_configs: set[SkinnyGemmConfig] = set() + for child in module.modules(): + is_linear = ( + isinstance(child, LinearBase) + and type(child.quant_method) is UnquantizedLinearMethod + ) + is_head = ( + isinstance(child, ParallelLMHead) + and type(child.quant_method) is UnquantizedEmbeddingMethod + ) + if not (is_linear or is_head): + continue + weight = getattr(child, "weight", None) + if weight is None or weight.dim() != 2: + continue + plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1])) + if plan is None: + continue + if is_linear: + child.quant_method = Qwen4ExpLowLatencyLinearMethod() + else: + child.quant_method = Qwen4ExpLowLatencyEmbeddingMethod() + warmup_configs.update(plan.values()) + + if warmup_configs: + shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs) diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py new file mode 100644 index 0000000000..de665095dc --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -0,0 +1,1112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Qwen4Exp model.""" + +from collections.abc import Iterable +from itertools import islice + +import torch +from torch import nn + +from vllm import envs +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + QwenGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateCopyFuncsByType, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import ( + HasInnerState, + IsHybrid, + MixtureOfExperts, + MultiModalEmbeddings, + SupportsLoRA, + SupportsMRoPE, + SupportsPP, + _require_is_multimodal, +) +from vllm.model_executor.models.qwen3_5 import ( + Qwen3_5ForConditionalGeneration, +) +from vllm.model_executor.models.qwen3_next import ( + Qwen3NextAttention, + Qwen3NextMLP, + Qwen3NextSparseMoeBlock, +) +from vllm.model_executor.models.qwen3_vl import ( + Qwen3_VisionTransformer, + Qwen3VLDummyInputsBuilder, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + StageMissingLayer, + WeightsMapper, + _merge_multimodal_embeddings, + extract_layer_index, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_fuse_shared_experts, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.sequence import IntermediateTensors +from vllm.tokenizers.registry import cached_tokenizer_from_config +from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, +) +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.kv_cache_interface import MambaSpec + +from ..config import Qwen4ExpConfig +from .hyperconnection import GatedResidual, HyperConnectionConfig + +try: + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm +except ModuleNotFoundError as exc: + # The Blackwell-only CuTe DSL helper is absent from the 1Cat SM70 tree. + if exc.name != "vllm.model_executor.kernels.linear.cute_dsl": + raise + + def enable_qwen4_exp_low_latency_gemm( + module: nn.Module, dtype: torch.dtype + ) -> None: + del module, dtype + + +from .ple_layer import Qwen4ExpPLELayer +from .qsa import Qwen4ExpQSAAttention + + +def without_modelopt_fp4( + quant_config: QuantizationConfig | None, +) -> QuantizationConfig | None: + """Return ``None`` for weights excluded from Qwen4Exp ModelOpt-FP4.""" + + if quant_config is not None and quant_config.get_name() == "modelopt_fp4": + return None + return quant_config + + +def _remap_qsa_cache_scale_name( + name: str, + qsa_layer_ids: frozenset[int], +) -> str: + """Map serialized main-cache scales onto the merged QSA owner. + + Regular attention keeps cache scales below its ``attn`` child. QSA owns + that cache directly, so only QSA layers need the final path component + moved to the owner's persistent ``_k_scale``/``_v_scale`` buffers. + """ + + scale_suffixes = { + "k_proj.k_scale": "_k_scale", + "k_proj.output_scale": "_k_scale", + "attn.k_scale": "_k_scale", + "attn._k_scale": "_k_scale", + "k_scale": "_k_scale", + "_k_scale": "_k_scale", + "v_proj.v_scale": "_v_scale", + "v_proj.output_scale": "_v_scale", + "attn.v_scale": "_v_scale", + "attn._v_scale": "_v_scale", + "v_scale": "_v_scale", + "_v_scale": "_v_scale", + } + for layer_id in qsa_layer_ids: + marker = f"layers.{layer_id}.self_attn." + marker_start = name.find(marker) + if marker_start < 0 or (marker_start > 0 and name[marker_start - 1] != "."): + continue + suffix = name[marker_start + len(marker) :] + mapped_suffix = scale_suffixes.get(suffix) + if mapped_suffix is not None: + return f"{name[: marker_start + len(marker)]}{mapped_suffix}" + return name + + +_QWEN4_EXP_IGNORED_MISSING_SUFFIXES = [ + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + "_weight_scale", + "_input_scale", +] + +# The checkpoint keeps down and injection projections separate; runtime packs +# them into adjacent logical shards of one MergedColumnParallelLinear. +_HC_WEIGHTS_MAPPER = WeightsMapper( + orig_to_new_stacked={ + "hyper_connection.input_mix_weight_down.weight": ( + "hyper_connection.input_mix_weight_down_block_inject.weight", + 0, + ), + "hyper_connection.block_inject_weight.weight": ( + "hyper_connection.input_mix_weight_down_block_inject.weight", + 1, + ), + } +) + +_QWEN3_5_WEIGHTS_MAPPER = WeightsMapper( + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".mlp.gate_proj": (".mlp.gate_up_proj", 0), + ".mlp.up_proj": (".mlp.gate_up_proj", 1), + ".shared_expert.gate_proj": (".shared_expert.gate_up_proj", 0), + ".shared_expert.up_proj": (".shared_expert.gate_up_proj", 1), + ".in_proj_qkv": (".in_proj_qkvz", (0, 1, 2)), + ".in_proj_z": (".in_proj_qkvz", 3), + ".in_proj_b": (".in_proj_ba", 0), + ".in_proj_a": (".in_proj_ba", 1), + } +) + + +class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock): + """Qwen3Next MoE with Qwen4Exp HC validation.""" + + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: + parallel_config = vllm_config.parallel_config + if parallel_config.use_sequence_parallel_moe: + raise NotImplementedError( + "Qwen4Exp HC does not support sequence-parallel MoE" + ) + super().__init__(vllm_config=vllm_config, prefix=prefix) + config = vllm_config.model_config.hf_text_config + self.n_shared_experts = int(config.shared_expert_intermediate_size > 0) + # Qwen3Next loads expert tensors manually in this private tree, so its + # FusedMoE does not receive the mapping required by AutoWeightsLoader. + # Qwen4Exp uses recursive loading and must attach that mapping here. + self.experts.expert_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.n_routed_experts, + num_redundant_experts=self.n_redundant_experts, + include_fused=True, + ) + + +class Qwen4ExpDecoderLayer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + layer_type: str, + prefix: str = "", + ) -> None: + super().__init__() + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + self.config = config + self.layer_type = layer_type + self.layer_idx = extract_layer_index(prefix) + if vllm_config.parallel_config.use_sequence_parallel_moe: + raise NotImplementedError( + "Qwen4Exp HC does not support sequence-parallel MoE" + ) + self.ple: Qwen4ExpPLELayer | None = None + ple_layer_ids = config.ple_layer_ids + if (self.layer_idx + 1) in ple_layer_ids: + ple_layer_ids_sorted = sorted(set(ple_layer_ids)) + ple_dense_layer_id_map = { + abs_id: idx for idx, abs_id in enumerate(ple_layer_ids_sorted) + } + ple_dense_layer_id = ple_dense_layer_id_map[self.layer_idx + 1] + self.ple = Qwen4ExpPLELayer( + config, + vllm_config=vllm_config, + layer_idx=self.layer_idx, + ple_dense_layer_id=ple_dense_layer_id, + prefix=f"{prefix}.ple", + ) + + if layer_type == "linear_attention": + self.linear_attn = QwenGatedDeltaNetAttention( + config, + vllm_config=vllm_config, + prefix=f"{prefix}.linear_attn", + gqa_interleaved_layout=False, + ) + elif layer_type == "full_attention": + use_qsa = getattr(config, "indexer_n_heads", None) is not None + if not use_qsa: + self.self_attn = Qwen3NextAttention( + config, + model_config=model_config, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + else: + self.self_attn = Qwen4ExpQSAAttention( + vllm_config=vllm_config, + config=config, + layer_id=self.layer_idx, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + else: + raise ValueError(f"Invalid layer_type {layer_type}") + + mlp_only_layers = getattr(config, "mlp_only_layers", []) + num_experts = getattr(config, "num_experts", 0) or 0 + absolute_layer_id = self.layer_idx + 1 + is_moe_layer = self.layer_idx not in mlp_only_layers and ( + num_experts > 0 and absolute_layer_id % config.decoder_sparse_step == 0 + ) + if is_moe_layer: + self.mlp = Qwen4ExpSparseMoeBlock( + vllm_config=vllm_config, prefix=f"{prefix}.mlp" + ) + else: + self.mlp = Qwen3NextMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + hc_config = HyperConnectionConfig( + hc_count=config.hc_count, + hidden_size=config.hidden_size, + params_dtype=model_config.dtype, + hc_lowrank=config.hc_lowrank, + rms_norm_eps=config.rms_norm_eps, + hc_per_branch_norm=True, + ) + self.attn_hyper_connection = GatedResidual( + hc_config, + prefix=maybe_prefix(prefix, "attn_hyper_connection"), + ) + self.mlp_hyper_connection = GatedResidual( + hc_config, + prefix=maybe_prefix(prefix, "mlp_hyper_connection"), + ) + + def forward( + self, + hidden_states: torch.Tensor, + prev_block_output: torch.Tensor | None, + prev_injection: torch.Tensor | None, + positions: torch.Tensor, + *, + input_ids: torch.Tensor | None, + query_start_loc: torch.Tensor | None, + ngram_context: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + attn_hc = self.attn_hyper_connection + if self.ple is not None: + # PLE adds directly to the multi-stream state, so pending HC state + # must be materialized before the addition. + if prev_block_output is not None and prev_injection is not None: + hidden_states = attn_hc.combine( + hidden_states, prev_block_output, prev_injection + ) + prev_block_output = prev_injection = None + + if input_ids is None or query_start_loc is None or ngram_context is None: + raise RuntimeError("PLE inputs were not prepared") + hidden_states = hidden_states + self.ple( + hidden_states, + input_ids, + query_start_loc, + ngram_context, + ) + + # Fuse a pending combine with this HC module's mix when possible. + if prev_block_output is not None and prev_injection is not None: + hidden_states, block_input, injection = attn_hc.combine_and_mix( + hidden_states, prev_block_output, prev_injection + ) + else: + hidden_states, block_input, injection = attn_hc.mix(hidden_states) + + if self.layer_type == "linear_attention": + use_direct_attention_output = ( + envs.VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM + and torch.compiler.is_compiling() + ) + attn_buffer = ( + None if use_direct_attention_output else torch.empty_like(block_input) + ) + projected_attn_out = self.linear_attn( + hidden_states=block_input, + output=attn_buffer, + ) + if use_direct_attention_output: + if projected_attn_out is None: + raise RuntimeError( + "SM70 TP4 fused prefill requires a direct GDN output" + ) + attn_out = projected_attn_out + else: + assert attn_buffer is not None + attn_out = attn_buffer + elif self.layer_type == "full_attention": + attn_out = self.self_attn( + hidden_states=block_input, + positions=positions, + output=None, + ) + else: + raise ValueError("Invalid layer_type") + + mlp_hc = self.mlp_hyper_connection + hidden_states, block_input, injection = mlp_hc.combine_and_mix( + hidden_states, attn_out, injection + ) + mlp_out = self.mlp(block_input) + return hidden_states, mlp_out, injection + + +class Qwen4ExpMixtureOfExperts(MixtureOfExperts): + """Expose Qwen4Exp routed experts through vLLM's EPLB protocol.""" + + def set_moe_parameters(self, layers: Iterable[nn.Module]) -> None: + self.moe_layers = [] + self.moe_mlp_layers = [] + example_moe = None + for layer in layers: + if isinstance(layer, Qwen4ExpDecoderLayer) and isinstance( + layer.mlp, Qwen4ExpSparseMoeBlock + ): + example_moe = layer.mlp + self.moe_mlp_layers.append(layer.mlp) + self.moe_layers.append(layer.mlp.experts) + + self.num_moe_layers = len(self.moe_layers) + if example_moe is None: + self.num_expert_groups = 0 + self.num_shared_experts = 0 + self.num_logical_experts = 0 + self.num_physical_experts = 0 + self.num_local_physical_experts = 0 + self.num_routed_experts = 0 + self.num_redundant_experts = 0 + return + + self.num_expert_groups = 1 + self.num_shared_experts = example_moe.n_shared_experts + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for moe in self.moe_mlp_layers: + moe.n_physical_experts = num_physical_experts + moe.n_local_physical_experts = num_local_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + "query_start_loc": 0, + "ngram_context": 0, + "deepstack_input_embeds": 0, + } +) +class Qwen4ExpModel(nn.Module): + hf_to_vllm_mapper = _QWEN3_5_WEIGHTS_MAPPER | _HC_WEIGHTS_MAPPER + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + self.config = config + self.num_redundant_experts = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) + self.vocab_size = config.vocab_size + self._qsa_layer_ids = frozenset( + layer_idx + for layer_idx, layer_type in enumerate(config.layer_types) + if layer_type == "full_attention" + and getattr(config, "indexer_n_heads", None) is not None + ) + self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size) + + def get_layer(prefix: str) -> Qwen4ExpDecoderLayer: + layer_idx = extract_layer_index(prefix) + return Qwen4ExpDecoderLayer( + vllm_config, + layer_type=config.layer_types[layer_idx], + prefix=prefix, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" + ) + intermediate_size = config.hidden_size * config.hc_count + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], intermediate_size + ) + + self.hyper_connection_mixer: GatedResidual | None + if get_pp_group().is_last_rank: + hc_config = HyperConnectionConfig( + hc_count=config.hc_count, + hidden_size=config.hidden_size, + params_dtype=vllm_config.model_config.dtype, + hc_lowrank=config.hc_lowrank, + rms_norm_eps=config.rms_norm_eps, + hc_per_branch_norm=True, + ) + self.hyper_connection_mixer = GatedResidual( + hc_config, + use_combine=False, + prefix=maybe_prefix(prefix, "hyper_connection_mixer"), + ) + else: + self.hyper_connection_mixer = None + + spec_config = vllm_config.speculative_config + # MTP HC multi-stream outputs: when speculative method=="mtp" and the + # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream + # hidden state [T, hc_count*H] so the MTP drafter can feed a real + # multi-stream backbone hidden on its first step (scheme A). Derived + # purely from config (NOT node identity) so P/D nodes stay consistent. + needs_mtp_hidden = ( + spec_config is not None + and getattr(spec_config, "method", None) == "mtp" + and get_pp_group().is_last_rank + ) + if needs_mtp_hidden: + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.hc_count * config.hidden_size, + dtype=vllm_config.model_config.dtype, + ) + else: + self._mtp_hidden_buffer = None + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + ngram_context: torch.Tensor | None = None, + deepstack_input_embeds: IntermediateTensors | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + if input_ids is None: + raise ValueError("input_ids or inputs_embeds is required") + hidden_states = self.embed_input_ids(input_ids) + hidden_states = hidden_states.repeat(1, self.config.hc_count) + else: + if intermediate_tensors is None: + raise ValueError("pipeline stage requires intermediate tensors") + hidden_states = intermediate_tensors["hidden_states"] + + block_output = None + injection = None + last_layer = None + for layer_idx, layer in islice( + enumerate(self.layers), self.start_layer, self.end_layer + ): + last_layer = layer + hidden_states, block_output, injection = layer( + hidden_states=hidden_states, + prev_block_output=block_output, + prev_injection=injection, + positions=positions, + input_ids=input_ids, + query_start_loc=query_start_loc, + ngram_context=ngram_context, + ) + if deepstack_input_embeds is not None and layer_idx < len( + deepstack_input_embeds + ): + deepstack_embed = deepstack_input_embeds[ + f"deepstack_input_embeds_{layer_idx}" + ] + deepstack_embed = ( + deepstack_embed.unsqueeze(-2) + .expand( + *deepstack_embed.shape[:-1], + self.config.hc_count, + self.config.hidden_size, + ) + .flatten(-2) + ) + # Deepstack is an external addition to the materialized + # multi-stream state and therefore terminates delayed combine. + hidden_states = layer.mlp_hyper_connection.combine( + hidden_states, block_output, injection + ) + block_output = None + injection = None + hidden_states = hidden_states + deepstack_embed + + if not get_pp_group().is_last_rank: + # PP transports one tensor, not the delayed HC tuple. Materialize + # with the HC module that produced the pending injection. + if last_layer is not None and block_output is not None: + hidden_states = last_layer.mlp_hyper_connection.combine( + hidden_states, block_output, injection + ) + return IntermediateTensors({"hidden_states": hidden_states}) + + # The final mixer consumes the last pending combine and returns both + # the sampled single stream and the materialized multi-stream state. + final_mixer = self.hyper_connection_mixer + assert final_mixer is not None + multi_hidden, sample_hidden_states, _ = final_mixer.combine_and_mix( + hidden_states, block_output, injection + ) + if self._mtp_hidden_buffer is not None: + # Capture the pre-final-mixer multi-stream hidden state + # [T, hc_count*H] for the MTP drafter (zero extra compute: + # this tensor is needed by the final mixer regardless). + num_tokens = multi_hidden.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_(multi_hidden) + return sample_hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = ( + ( + _remap_qsa_cache_scale_name(name, self._qsa_layer_ids), + weight, + ) + for name, weight in weights + ) + weights = maybe_fuse_shared_experts( + weights, + n_routed_experts=getattr(self.config, "num_experts", 0) or 0, + n_shared_experts=1, + ckpt_prefix="mlp.shared_expert", + ) + # Non-persistent PLE state rebuilt in __init__; skip any ckpt + # column for them. + skip_substrs = [ + "hashstats_", + "token_lookup", + "hyper_connection_mixer.block_inject_weight", + ] + loader = AutoWeightsLoader( + self, + skip_substrs=skip_substrs, + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) + loaded = loader.load_weights( + weights, + mapper=self.hf_to_vllm_mapper, + ) + return loaded + + +class Qwen4ExpForCausalLM( + nn.Module, + HasInnerState, + SupportsLoRA, + SupportsMRoPE, + SupportsPP, + Qwen4ExpMixtureOfExperts, + IsHybrid, +): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], + "in_proj_ba": ["in_proj_b", "in_proj_a"], + "input_mix_weight_down_block_inject": [ + "input_mix_weight_down", + "block_inject_weight", + "_input_mix_padding", + ], + } + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={"model.language_model.": "model."} + ) + requires_raw_input_tokens = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.quant_config = vllm_config.quant_config + self.config = config + self.scheduler_config = vllm_config.scheduler_config + if vllm_config.cache_config.mamba_cache_mode == "all": + raise NotImplementedError( + "Qwen4Exp currently does not support 'all' prefix caching, " + "please use '--mamba-cache-mode=align' instead" + ) + self.model = Qwen4ExpModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) + self.set_moe_parameters(self.model.layers) + enable_qwen4_exp_low_latency_gemm(self, self.model_config.dtype) + + @staticmethod + def get_model_state_cls(): + from .model_state import Qwen4ExpModelState + + return Qwen4ExpModelState + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + # Forward kwargs unchanged so the runner's _maybe_add_ngram_kwargs + # path (query_start_loc / ngram_context) reaches Qwen4ExpModel. + return self.model( + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + **kwargs, + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + @classmethod + def get_ple_mamba_state_dtype_from_config( + cls, + vllm_config: VllmConfig, + ) -> tuple[torch.dtype, ...]: + return MambaStateDtypeCalculator.short_conv_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + ) + + @classmethod + def get_ple_mamba_state_shape_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[tuple[int, int]]: + hf_config = vllm_config.model_config.hf_text_config + conv_kernel_size = hf_config.ple_conv_kernel_size + short_conv_dilation = hf_config.ngram_size + conv_state_len = (conv_kernel_size - 1) * short_conv_dilation + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + hc_count = hf_config.hc_count + hc_hidden_size = hf_config.hidden_size * hc_count + return MambaStateShapeCalculator.short_conv_state_shape( + tp_world_size=1, + intermediate_size=hc_hidden_size, + conv_kernel=conv_state_len + 1, + num_spec=num_spec, + ) + + @classmethod + def get_gdn_mamba_state_dtype_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.gated_delta_net_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, + ) + + @classmethod + def get_gdn_mamba_state_shape_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_text_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.gated_delta_net_state_shape( + tp_size, + hf_config.linear_num_key_heads, + hf_config.linear_num_value_heads, + hf_config.linear_key_head_dim, + hf_config.linear_value_head_dim, + hf_config.linear_conv_kernel_dim, + num_spec, + ) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: VllmConfig, + ) -> tuple[torch.dtype, torch.dtype]: + return cls.get_gdn_mamba_state_dtype_from_config(vllm_config) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + return cls.get_gdn_mamba_state_shape_from_config(vllm_config) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func() + + @classmethod + def get_mamba_state_copy_funcs( + cls, + mamba_types: set[MambaAttentionBackendEnum], + ) -> MambaStateCopyFuncsByType: + copy_funcs_by_type = { + MambaAttentionBackendEnum.GDN_ATTN: cls.get_mamba_state_copy_func(), + MambaAttentionBackendEnum.SHORT_CONV: ( + MambaStateCopyFuncCalculator.short_conv_state_copy_func() + ), + } + missing_types = mamba_types - copy_funcs_by_type.keys() + assert not missing_types, f"missing state copy funcs for {missing_types}" + return { + mamba_type: copy_funcs_by_type[mamba_type] for mamba_type in mamba_types + } + + @classmethod + def get_mamba_specs_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[MambaSpec, ...]: + """Return all MambaSpecs for this model (GDN layers + PLE layer). + + The PLE layer uses a separate short_conv MambaSpec whose page_size_bytes + may exceed the GDN spec; callers should take the maximum. + """ + return ( + MambaSpec( + shapes=cls.get_gdn_mamba_state_shape_from_config(vllm_config), + dtypes=cls.get_gdn_mamba_state_dtype_from_config(vllm_config), + block_size=-1, + ), + MambaSpec( + shapes=cls.get_ple_mamba_state_shape_from_config(vllm_config), + dtypes=cls.get_ple_mamba_state_dtype_from_config(vllm_config), + block_size=-1, + tp_replicated=True, + ), + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + return self.model._mtp_hidden_buffer + + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec], + ) -> tuple[torch.Tensor, int]: + positions = torch.arange(len(input_tokens), dtype=torch.long) + return positions.unsqueeze(0).expand(3, -1), 0 + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_substrs=["mtp."], + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +class Qwen4ExpProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_config(self) -> Qwen4ExpConfig: + return self.ctx.get_hf_config(Qwen4ExpConfig) + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Qwen4ExpProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class Qwen4ExpForConditionalGeneration( + Qwen3_5ForConditionalGeneration, + HasInnerState, + Qwen4ExpMixtureOfExperts, +): + """Qwen3-VL vision tower backed by the Qwen4Exp language model.""" + + requires_raw_input_tokens = True + + packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | { + "input_mix_weight_down_block_inject": [ + "input_mix_weight_down", + "block_inject_weight", + "_input_mix_padding", + ] + } + + @staticmethod + def get_model_state_cls(): + from .model_state import Qwen4ExpModelState + + return Qwen4ExpModelState + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None: + nn.Module.__init__(self) + config: Qwen4ExpConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + if multimodal_config is None: + raise ValueError( + "Qwen4ExpForConditionalGeneration requires multimodal_config" + ) + + self.config = config + self.model_config = vllm_config.model_config + self.multimodal_config = multimodal_config + self.language_model_only = multimodal_config.language_model_only + if self.language_model_only: + self.use_data_parallel = False + self.is_multimodal_pruning_enabled = False + self.video_pruning_method = None + self.video_pruning_rate = 0.0 + self._tokenizer = None + self.visual = StageMissingLayer("vision_tower") + self._tower_model_names = [] + else: + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + self._init_video_pruning(multimodal_config) + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) + + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.visual = Qwen3_VisionTransformer( + config.vision_config, + norm_eps=config.text_config.rms_norm_eps, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "visual"), + ) + + self.use_deepstack = ( + not self.language_model_only + and bool(config.vision_config.deepstack_visual_indexes) + and not isinstance(self.visual, StageMissingLayer) + ) + self.deepstack_num_level = ( + len(config.vision_config.deepstack_visual_indexes) + if self.use_deepstack + else 0 + ) + self.visual_dim = config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level + + if self.use_deepstack: + self.deepstack_input_embeds = [ + torch.zeros( + vllm_config.scheduler_config.max_num_batched_tokens, + config.text_config.hidden_size, + ) + for _ in range(self.deepstack_num_level) + ] + self.deepstack_input_embeds_num_tokens = 0 + + with self._mark_language_model(vllm_config): + self.language_model = Qwen4ExpForCausalLM( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "language_model"), + ) + + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + if not get_pp_group().is_first_rank and self.use_deepstack: + assert self.language_model.model.start_layer >= len( + config.vision_config.deepstack_visual_indexes + ), ( + "start_layer should be greater than or equal to " + "len(deepstack_visual_indexes)" + ) + self.set_moe_parameters(self.language_model.model.layers) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self._embed_text_input_ids( + input_ids, + self.language_model.embed_input_ids, + is_multimodal=is_multimodal, + ) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + if self.language_model_only: + raise ValueError( + "Qwen4Exp language_model_only does not accept multimodal embeddings" + ) + + is_multimodal = _require_is_multimodal(is_multimodal) + if self.use_deepstack: + deepstack_input_embeds, multimodal_embeddings = ( + self._compute_deepstack_embeds( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + ) + else: + deepstack_input_embeds = None + + inputs_embeds = _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + if deepstack_input_embeds is not None: + self._set_deepstack_input_embeds(deepstack_input_embeds) + return inputs_embeds + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + return self.language_model.get_mtp_target_hidden_states() + + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.language_model.get_top_tokens(hidden_states) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + if inputs_embeds is not None and get_pp_group().is_first_rank: + deepstack_input_embeds = self._get_deepstack_input_embeds( + inputs_embeds.size(0) + ) + else: + deepstack_input_embeds = None + + hidden_states = self.language_model.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + query_start_loc=kwargs.get("query_start_loc"), + ngram_context=kwargs.get("ngram_context"), + deepstack_input_embeds=deepstack_input_embeds, + ) + if inputs_embeds is not None and get_pp_group().is_first_rank: + self._clear_deepstack_input_embeds(inputs_embeds.size(0)) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=["visual."] if self.language_model_only else None, + skip_substrs=["mtp."], + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: VllmConfig, + ) -> tuple[torch.dtype, torch.dtype]: + return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config) + + @classmethod + def get_mamba_state_shape_from_config( # type: ignore[override] + cls, + vllm_config: VllmConfig, + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return Qwen4ExpForCausalLM.get_mamba_state_copy_func() + + @classmethod + def get_mamba_state_copy_funcs( + cls, + mamba_types: set[MambaAttentionBackendEnum], + ) -> MambaStateCopyFuncsByType: + return Qwen4ExpForCausalLM.get_mamba_state_copy_funcs(mamba_types) + + @classmethod + def get_mamba_specs_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[MambaSpec, ...]: + return Qwen4ExpForCausalLM.get_mamba_specs_from_config(vllm_config) + + +__all__ = [ + "Qwen4ExpDecoderLayer", + "Qwen4ExpForCausalLM", + "Qwen4ExpForConditionalGeneration", + "Qwen4ExpMixtureOfExperts", + "Qwen4ExpModel", + "Qwen4ExpSparseMoeBlock", +] diff --git a/vllm/models/qwen4_exp/nvidia/model_state.py b/vllm/models/qwen4_exp/nvidia/model_state.py new file mode 100644 index 0000000000..ec2a0fef8b --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/model_state.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Model-runner state for Qwen4Exp PLE inputs.""" + +from typing import Any + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState +from vllm.v1.worker.gpu.states import RequestState + + +class Qwen4ExpModelState(MambaHybridModelState): + """Add rollback-safe PLE n-gram context to the model inputs.""" + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, + ) -> None: + super().__init__(vllm_config, model, encoder_cache, device) + config = self.model_config.hf_text_config + self.uses_ngram_embedding = bool(config.ple_layer_ids) + if not self.uses_ngram_embedding: + self.ngram_context_len = 0 + self.ngram_eos_token_id = 0 + return + + if vllm_config.parallel_config.pipeline_parallel_size > 1: + raise RuntimeError( + "N-gram PLE embedding currently requires " + "pipeline_parallel_size=1 because non-first pipeline ranks do " + "not receive the raw input_ids required by PLE. Please run " + "with PP=1." + ) + + self.ngram_context_len = int(config.ngram_size) - 1 + if self.ngram_context_len <= 0: + raise ValueError("N-gram embedding requires context length >= 1.") + self.ngram_eos_token_id = int(config.eos_token_id) + self.ngram_context = torch.full( + (self.max_num_reqs, self.ngram_context_len), + self.ngram_eos_token_id, + dtype=torch.int32, + device=self.device, + ) + self.ngram_context_offsets = torch.arange( + -self.ngram_context_len, + 0, + dtype=torch.int64, + device=self.device, + ) + self.ple_query_start_loc = torch.zeros( + self.max_num_reqs + 1, + dtype=torch.int32, + device=self.device, + ) + + def _prepare_ngram_context( + self, + input_batch: InputBatch, + req_states: RequestState, + ) -> torch.Tensor: + num_reqs = input_batch.num_reqs + num_reqs_padded = input_batch.num_reqs_after_padding + context = self.ngram_context[:num_reqs_padded] + context.fill_(self.ngram_eos_token_id) + if num_reqs == 0: + return context + + request_indices = input_batch.idx_mapping[:num_reqs].long() + context_end = req_states.num_computed_tokens.gpu[request_indices].long() + token_indices = context_end.unsqueeze(1) + self.ngram_context_offsets + valid_tokens = token_indices >= 0 + token_indices.clamp_min_(0) + context_tokens = req_states.all_token_ids.gpu[ + request_indices.unsqueeze(1), token_indices + ] + context[:num_reqs].copy_( + torch.where( + valid_tokens, + context_tokens, + context_tokens.new_full((), self.ngram_eos_token_id), + ) + ) + return context + + def prepare_inputs( + self, + input_batch: InputBatch, + req_states: RequestState, + ) -> dict[str, Any]: + model_inputs = super().prepare_inputs(input_batch, req_states) + if not self.uses_ngram_embedding: + return model_inputs + + num_reqs_padded = input_batch.num_reqs_after_padding + query_start_loc = self.ple_query_start_loc[: num_reqs_padded + 1] + query_start_loc.copy_(input_batch.query_start_loc[: num_reqs_padded + 1]) + model_inputs.update( + query_start_loc=query_start_loc, + ngram_context=self._prepare_ngram_context(input_batch, req_states), + ) + return model_inputs + + def prepare_dummy_inputs( + self, + num_reqs: int, + num_tokens: int, + ) -> dict[str, Any]: + model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens) + if not self.uses_ngram_embedding: + return model_inputs + + query_start_loc = self.ple_query_start_loc[: num_reqs + 1] + query_start_loc[0] = 0 + tokens_per_req, num_extra_tokens = divmod(num_tokens, num_reqs) + query_lens = torch.full( + (num_reqs,), + tokens_per_req, + dtype=query_start_loc.dtype, + device=query_start_loc.device, + ) + if num_extra_tokens > 0: + query_lens[-num_extra_tokens:] += 1 + torch.cumsum(query_lens, dim=0, out=query_start_loc[1:]) + + ngram_context = self.ngram_context[:num_reqs] + ngram_context.fill_(self.ngram_eos_token_id) + model_inputs.update( + query_start_loc=query_start_loc, + ngram_context=ngram_context, + ) + return model_inputs + + +__all__ = ["Qwen4ExpModelState"] diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py new file mode 100644 index 0000000000..18f4575f54 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -0,0 +1,489 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Qwen4Exp MTP (Multi-Token Predictor) model. + +The MTP draft model reuses the Qwen4Exp backbone (PLE/HC/MoE) but: + - drops all multi-modal handling (text-only), + - forces PLE off while keeping the main model's HC stream count, + - fuses the backbone hidden and the new-token embedding via + ``residual_linear_shared`` (fc_embedding + shared fc_hidden) instead of + the ``Linear(2H, H)`` + repeat used by other MTP variants, + - emits TWO hidden streams per step (scheme A): a single stream [T, H] + (final-mixer collapsed, fed to the LM head) and a pre-final-mixer + multi stream [T, hc_count*H] (fed to the next draft step). +""" + +from collections.abc import Iterable + +import regex as re +import torch +from torch import nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig, replace, set_current_vllm_config +from vllm.distributed import get_pp_group +from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.utils import configure_quant_config +from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + get_draft_quant_config, + make_empty_intermediate_tensors_factory, + maybe_fuse_shared_experts, + maybe_prefix, +) +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, +) + +from .hyperconnection import GatedResidual, HyperConnectionConfig + +try: + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm +except ModuleNotFoundError as exc: + # The Blackwell-only CuTe DSL helper is absent from the 1Cat SM70 tree. + if exc.name != "vllm.model_executor.kernels.linear.cute_dsl": + raise + + def enable_qwen4_exp_low_latency_gemm( + module: nn.Module, dtype: torch.dtype + ) -> None: + del module, dtype + + +from .model import ( + _HC_WEIGHTS_MAPPER, + _QWEN3_5_WEIGHTS_MAPPER, + _QWEN4_EXP_IGNORED_MISSING_SUFFIXES, + Qwen4ExpDecoderLayer, + Qwen4ExpMixtureOfExperts, +) + + +def _remap_ignored_layers( + ignored_layers: list[str], + mtp_start_layer_idx: int, +) -> list[str]: + remapped: list[str] = [] + for name in ignored_layers: + if name.startswith("mtp."): + new_name = re.sub( + r"(?<=\.layers\.)\d+", + lambda m: str(mtp_start_layer_idx + int(m.group(0))), + name, + ) + remapped.append(new_name) + else: + remapped.append(name) + return remapped + + +def _remap_mtp_weight_name(name: str) -> str | None: + """Map Qwen4Exp checkpoint paths into the standalone draft model.""" + + for checkpoint_prefix in ( + "model.language_model.", + "language_model.", + ): + if name.startswith(checkpoint_prefix): + name = name.removeprefix(checkpoint_prefix) + break + + if name.startswith("embed_tokens."): + name = f"model.{name}" + if name.startswith("model.mtp."): + name = name.removeprefix("model.") + if name.startswith("mtp.shared_head.head."): + return name.replace("mtp.shared_head.head.", "lm_head.", 1) + if name.startswith("model.shared_head.head."): + return name.replace("model.shared_head.head.", "lm_head.", 1) + if name.startswith("shared_head.head."): + return name.replace("shared_head.head.", "lm_head.", 1) + if name.startswith("model.lm_head."): + return name.removeprefix("model.") + if name.startswith("mtp."): + return name.replace("mtp.", "model.", 1) + if name.startswith("model.embed_tokens.") or name.startswith("lm_head."): + return name + return None + + +def _validate_mtp_expert_weights_loaded( + model: nn.Module, + loaded_weights: set[str], +) -> None: + """Reject a silently incomplete Qwen4Exp MTP routed-expert load.""" + + required_suffixes = ( + ".mlp.experts.w13_weight", + ".mlp.experts.w2_weight", + ) + required = { + name for name, _ in model.named_parameters() if name.endswith(required_suffixes) + } + missing = required - loaded_weights + if missing: + missing_names = ", ".join(sorted(missing)) + raise ValueError( + "Qwen4Exp MTP routed-expert checkpoint weights were not loaded: " + f"{missing_names}. Check fused/per-expert checkpoint mappings." + ) + + +def _make_draft_vllm_config( + vllm_config: VllmConfig, + mtp_start_layer_idx: int, +) -> VllmConfig: + """Ensure that the draft model config is set in the vLLM config.""" + speculative_config = vllm_config.speculative_config + if speculative_config is None or speculative_config.draft_model_config is None: + raise ValueError("speculative_config.draft_model_config must be set") + + draft_quant_config = get_draft_quant_config(vllm_config) + + # inject packed and ignored modules to the quantization config of draft model + if draft_quant_config is not None: + configure_quant_config(draft_quant_config, Qwen4ExpMTP) + ignored_layers = getattr(draft_quant_config, "ignored_layers", None) + if ignored_layers: + setattr( # noqa: B010 + draft_quant_config, + "ignored_layers", + _remap_ignored_layers(ignored_layers, mtp_start_layer_idx), + ) + exclude_modules = getattr(draft_quant_config, "exclude_modules", None) + if exclude_modules: + setattr( # noqa: B010 + draft_quant_config, + "exclude_modules", + _remap_ignored_layers(exclude_modules, mtp_start_layer_idx), + ) + + draft_vllm_config = replace( + vllm_config, + model_config=speculative_config.draft_model_config, + ) + # VllmConfig post-init derives the target quant config, so restore the + # independently resolved draft quant config after replacement. + draft_vllm_config.quant_config = draft_quant_config + return draft_vllm_config + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + "hidden_states": 0, + } +) +class Qwen4ExpMultiTokenPredictor(nn.Module): + hf_to_vllm_mapper = _QWEN3_5_WEIGHTS_MAPPER | _HC_WEIGHTS_MAPPER + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + model_config = vllm_config.model_config + config: Qwen4ExpTextConfig = model_config.hf_text_config + + self.config = config + self.vocab_size = config.vocab_size + + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1) + + self.hidden_size = config.hidden_size + self.hc_count = config.hc_count + + self.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size) + draft_vllm_config = _make_draft_vllm_config( + vllm_config, + self.mtp_start_layer_idx, + ) + with set_current_vllm_config(draft_vllm_config, prefix=prefix): + # residual_linear_shared fusion: fc_embedding projects the token + # embedding, fc_hidden (shared across HC branches) projects the + # backbone hidden; the embedding is added as a residual to every + # branch (see mtp_residual_linear_shared.md). + self.fc_embedding = ColumnParallelLinear( + self.hidden_size, + self.hidden_size, + gather_output=True, + bias=False, + return_bias=False, + quant_config=draft_vllm_config.quant_config, + prefix=f"{prefix}.fc_embedding", + ) + self.fc_hidden = ColumnParallelLinear( + self.hidden_size, + self.hidden_size, + gather_output=True, + bias=False, + return_bias=False, + quant_config=draft_vllm_config.quant_config, + prefix=f"{prefix}.fc_hidden", + ) + self.layers = nn.ModuleList( + Qwen4ExpDecoderLayer( + draft_vllm_config, + layer_type="full_attention", + prefix=f"{prefix}.layers.{self.mtp_start_layer_idx + idx}", + ) + for idx in range(self.num_mtp_layers) + ) + + self.pre_fc_norm_embedding = GemmaRMSNorm( + self.hidden_size, eps=config.rms_norm_eps + ) + self.pre_fc_norm_hidden = GemmaRMSNorm( + self.hidden_size * self.hc_count, eps=config.rms_norm_eps + ) + # HC final mixer collapses the multi stream into [T, H] for the LM head. + hc_config = HyperConnectionConfig( + hc_count=config.hc_count, + hidden_size=config.hidden_size, + params_dtype=model_config.dtype, + hc_lowrank=config.hc_lowrank, + rms_norm_eps=config.rms_norm_eps, + hc_per_branch_norm=True, + ) + self.hyper_connection_mixer = GatedResidual( + hc_config, + use_combine=False, + prefix=maybe_prefix(prefix, "hyper_connection_mixer"), + ) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], self.hidden_size * self.hc_count + ) + + def _iter_qsa_attentions(self): + """Yield MTP attention modules that own a QSA indexer.""" + + for layer in self.layers: + attention = getattr(layer, "self_attn", None) + if ( + attention is not None + and getattr(attention, "indexer", None) is not None + ): + yield attention + + def set_skip_topk(self, skip: bool) -> None: + """Select on MTP step 0 and reuse its QSA indices on later steps.""" + + for attention in self._iter_qsa_attentions(): + attention.indexer.skip_topk = skip + + def compact_topk_indices(self, row_indices: torch.Tensor) -> None: + """Keep each request's target-aligned step-0 sparse-index row.""" + + num_rows = row_indices.numel() + for attention in self._iter_qsa_attentions(): + buffer = attention.topk_indices_buffer + selected = buffer.index_select(0, row_indices) + buffer[:num_rows].copy_(selected) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor | None = None, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors: + hc_count = self.hc_count + hidden_size = self.hidden_size + + if get_pp_group().is_first_rank: + assert hidden_states is not None + if inputs_embeds is None: + assert input_ids is not None + inputs_embeds = self.embed_input_ids(input_ids) + # Embedding branch: pre-norm -> fc_embedding -> [T, H]. + inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds) + inputs_embeds = self.fc_embedding(inputs_embeds) + + # Backbone hidden is multi-stream [T, hc_count*H] (scheme A: + # the main model truly emits the pre-final-mixer multi stream + # on the first step; subsequent steps reuse the prior draft + # step's multi stream). + num_tokens = hidden_states.shape[0] + hidden_states = hidden_states.view(num_tokens, hc_count, hidden_size) + hidden_states = self.pre_fc_norm_hidden(hidden_states.flatten(-2)).view( + num_tokens, hc_count, hidden_size + ) + hidden_states = self.fc_hidden(hidden_states) + # Add the embedding residual to every branch, then fold back + # to [T, hc_count*H] (HC outer, HS inner) for the HC decoder. + hidden_states = inputs_embeds.unsqueeze(-2) + hidden_states + hidden_states = hidden_states.flatten(-2) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + current_step_idx = spec_step_idx % self.num_mtp_layers + layer = self.layers[current_step_idx] + hidden_states, block_output, injection = layer( + hidden_states=hidden_states, + prev_block_output=None, + prev_injection=None, + positions=positions, + input_ids=None, + query_start_loc=None, + ngram_context=None, + ) + if not get_pp_group().is_last_rank: + # As in the target model, PP carries a materialized tensor rather + # than the delayed hidden/output/injection tuple. + hidden_states = layer.mlp_hyper_connection.combine( + hidden_states, block_output, injection + ) + return IntermediateTensors({"hidden_states": hidden_states}) + + # Last PP rank finalize. Keep both: + # (A) sample_hidden_states [T, H] -> single stream for the LM head + # (B) multi_hidden [T, hc_count*H] -> pre-final-mixer multi stream + # for the next draft step (zero extra compute, just kept). + multi_hidden, sample_hidden_states, _ = ( + self.hyper_connection_mixer.combine_and_mix( + hidden_states, block_output, injection + ) + ) + return sample_hidden_states, multi_hidden + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = maybe_fuse_shared_experts( + weights, + n_routed_experts=getattr(self.config, "num_experts", 0) or 0, + n_shared_experts=1, + ckpt_prefix="mlp.shared_expert", + ) + loader = AutoWeightsLoader( + self, + skip_substrs=["hyper_connection_mixer.block_inject_weight"], + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + "hidden_states": 0, + } +) +class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], + "in_proj_ba": ["in_proj_b", "in_proj_a"], + "input_mix_weight_down_block_inject": [ + "input_mix_weight_down", + "block_inject_weight", + "_input_mix_padding", + ], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + cache_config = vllm_config.cache_config + if cache_config.mamba_cache_mode == "all": + raise NotImplementedError( + "Qwen4ExpMTP currently does not support 'all' prefix caching, " + "please use '--mamba-cache-mode=align' instead" + ) + + self.quant_config = vllm_config.quant_config + + super().__init__() + self.config = config + self.model = Qwen4ExpMultiTokenPredictor( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "mtp"), + ) + + if get_pp_group().is_last_rank: + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) + self.set_moe_parameters(self.model.layers) + enable_qwen4_exp_low_latency_gemm(self, vllm_config.model_config.dtype) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( # type: ignore[override] + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor | None = None, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors: + return self.model( + input_ids, + positions, + hidden_states, + intermediate_tensors, + inputs_embeds, + spec_step_idx=spec_step_idx, + ) + + def compute_logits( + self, hidden_states: torch.Tensor, spec_step_idx: int = 0 + ) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def get_top_tokens( + self, hidden_states: torch.Tensor, spec_step_idx: int = 0 + ) -> torch.Tensor: + return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + def remap_weight_names(): + for name, weight in weights: + remapped_name = _remap_mtp_weight_name(name) + if remapped_name is not None: + yield remapped_name, weight + + loader = AutoWeightsLoader( + self, + skip_substrs=["hyper_connection_mixer.block_inject_weight"], + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) + loaded_weights = loader.load_weights(remap_weight_names()) + _validate_mtp_expert_weights_loaded(self, loaded_weights) + return loaded_weights + + +__all__ = ["Qwen4ExpMTP", "Qwen4ExpMultiTokenPredictor"] diff --git a/vllm/models/qwen4_exp/nvidia/ops/__init__.py b/vllm/models/qwen4_exp/nvidia/ops/__init__.py new file mode 100644 index 0000000000..779c57191d --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA-only Qwen4Exp kernels; import leaf modules directly.""" diff --git a/vllm/models/qwen4_exp/nvidia/ops/hc.py b/vllm/models/qwen4_exp/nvidia/ops/hc.py new file mode 100644 index 0000000000..cde145f6fb --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/hc.py @@ -0,0 +1,491 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA HyperConnection kernels for Qwen4Exp.""" + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import direct_register_custom_op + + +@triton.jit +def _grouped_gemma_rmsnorm_kernel( + x_ptr, + w_ptr, + y_ptr, + stride_x, + stride_y, + DIM: tl.constexpr, + NUM_GROUPS: tl.constexpr, + W_SHARED: tl.constexpr, + EPS: tl.constexpr, + launch_pdl: tl.constexpr, +) -> None: + GROUP_DIM: tl.constexpr = DIM // NUM_GROUPS + BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(GROUP_DIM) + + pid = tl.program_id(0) + group_id = pid % NUM_GROUPS + row = pid // NUM_GROUPS + + offs_g = tl.arange(0, BLOCK_SIZE) + offsets = group_id * GROUP_DIM + offs_g + mask = offs_g < GROUP_DIM + # A [GROUP_DIM] affine is shared; a [DIM] affine follows the grouped + # checkpoint layout. + w_offs = offs_g if W_SHARED else offsets + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + x = tl.load(x_ptr + row * stride_x + offsets, mask, other=0.0).to(tl.float32) + w = tl.load(w_ptr + w_offs, mask, other=0.0) + + rrms = tl.rsqrt(tl.sum(x * x) / GROUP_DIM + EPS) + # Gemma's (1 + w) affine is written this way to lower to an FMA. + y = x * rrms + y += y * w.to(tl.float32) + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + tl.store(y_ptr + row * stride_y + offsets, y, mask) + + +def _grouped_gemma_rmsnorm( + x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int +) -> torch.Tensor: + N, DIM = x.shape + assert x.stride(1) == 1, "grouped Gemma RMSNorm requires unit inner stride" + assert weight.is_contiguous(), "grouped Gemma RMSNorm weight must be contiguous" + assert DIM % num_groups == 0 + group_dim = DIM // num_groups + assert weight.numel() in (group_dim, DIM) + + y = x.new_empty(x.shape) + _grouped_gemma_rmsnorm_kernel[(N * num_groups,)]( + x, + weight, + y, + x.stride(0), + y.stride(0), + DIM, + num_groups, + W_SHARED=weight.numel() == group_dim, + EPS=eps, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return y + + +@triton.jit +def _hc_silu_kernel( + x_ptr, + y_ptr, + stride_x, + stride_y, + DIM: tl.constexpr, + HC: tl.constexpr, + launch_pdl: tl.constexpr, +) -> None: + BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(DIM) + + row = tl.program_id(0) + offs = tl.arange(0, BLOCK_SIZE) + mask = offs < DIM + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + x = tl.load(x_ptr + row * stride_x + offs, mask).to(tl.float32) / HC + y = x * tl.sigmoid(x) + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + tl.store(y_ptr + row * stride_y + offs, y, mask) + + +def _hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor: + num_tokens, DIM = x.shape + assert x.stride(1) == 1 + + output = x.new_empty(x.shape) + _hc_silu_kernel[(num_tokens,)]( + x, + output, + x.stride(0), + output.stride(0), + DIM=DIM, + HC=hc_count, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return output + + +@triton.jit +def _hc_gate_mix_kernel( + x_ptr, + g_ptr, + y_ptr, + stride_x, + stride_g, + stride_y, + DIM: tl.constexpr, + HC: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +) -> None: + HC_DIM: tl.constexpr = DIM // HC + + row = tl.program_id(0) + tile_id = tl.program_id(1) + offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs_inner < HC_DIM + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + # The constexpr loop is unrolled and keeps one stream live at a time. + # Materializing [HC, BLOCK_SIZE] more than doubles latency at large M. + acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for stream in tl.static_range(HC): + offsets = stream * HC_DIM + offs_inner + g = tl.load(g_ptr + row * stride_g + offsets, mask, other=0.0) + x = tl.load(x_ptr + row * stride_x + offsets, mask, other=0.0) + acc += tl.sigmoid(g.to(tl.float32)) * x.to(tl.float32) + acc /= HC + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + tl.store(y_ptr + row * stride_y + offs_inner, acc, mask) + + +def _hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor: + N, DIM = gate.shape + assert x.shape == gate.shape + assert DIM % hc_count == 0 + assert x.stride(1) == 1 + assert gate.stride(1) == 1 + + HC_DIM = DIM // hc_count + out = x.new_empty(N, HC_DIM) + BLOCK_SIZE = 512 + _hc_gate_mix_kernel[(N, triton.cdiv(HC_DIM, BLOCK_SIZE))]( + x, + gate, + out, + x.stride(0), + gate.stride(0), + out.stride(0), + DIM, + hc_count, + BLOCK_SIZE, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out + + +@triton.jit +def _hc_combine_kernel( + block_ptr, + res_ptr, + inj_ptr, + out_ptr, + stride_block, + stride_res, + stride_inj, + stride_out, + HC_DIM: tl.constexpr, + HC: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +) -> None: + HC_PAD: tl.constexpr = triton.next_power_of_2(HC) + + row = tl.program_id(0) + tile_id = tl.program_id(1) + + offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask_inner = offs_inner < HC_DIM + offs_hc = tl.arange(0, HC_PAD) + mask_hc = offs_hc < HC + offs = offs_hc[:, None] * HC_DIM + offs_inner[None, :] + mask = mask_hc[:, None] & mask_inner[None, :] + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0) + block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0) + res = tl.load(res_ptr + row * stride_res + offs, mask, other=0.0) + + # Keeping HC as a broadcast dimension is faster here than four separate + # residual load/store sequences. + inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC) + out = res.to(tl.float32) + block.to(tl.float32)[None, :] * inj[:, None] + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + tl.store(out_ptr + row * stride_out + offs, out, mask=mask) + + +def _hc_combine( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + hc_count: int, +) -> torch.Tensor: + N, DIM = residual.shape + assert DIM % hc_count == 0 + hc_dim = DIM // hc_count + assert block_output.shape == (N, hc_dim) + assert injection_logits.shape == (N, hc_count) + assert residual.stride(1) == 1 + assert block_output.stride(1) == 1 + assert injection_logits.stride(1) == 1 + + out = residual.new_empty(residual.shape) + BLOCK_SIZE = 512 + _hc_combine_kernel[(N, triton.cdiv(hc_dim, BLOCK_SIZE))]( + block_output, + residual, + injection_logits, + out, + block_output.stride(0), + residual.stride(0), + injection_logits.stride(0), + out.stride(0), + hc_dim, + hc_count, + BLOCK_SIZE, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out + + +@triton.jit +def _hc_combine_norm_kernel( + block_ptr, + res_ptr, + inj_ptr, + w_ptr, + out_ptr, + y_ptr, + stride_block, + stride_res, + stride_inj, + stride_out, + stride_y, + HC_DIM: tl.constexpr, + HC: tl.constexpr, + W_SHARED: tl.constexpr, + EPS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +) -> None: + HC_PAD: tl.constexpr = triton.next_power_of_2(HC) + NUM_TILES: tl.constexpr = triton.cdiv(HC_DIM, BLOCK_SIZE) + NUM_TILES_PAD: tl.constexpr = triton.next_power_of_2(NUM_TILES) + + row = tl.program_id(0) + stream = tl.program_id(1) + offs_hc = tl.arange(0, HC_PAD) + mask_hc = offs_hc < HC + tile_ids = tl.arange(0, NUM_TILES_PAD) + offs_inner = tile_ids[:, None] * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :] + mask_inner = offs_inner < HC_DIM + offs = stream * HC_DIM + offs_inner + # Shared norm weights repeat across streams; per-branch weights use the + # same flattened HC layout as the residual. + w_offs = offs_inner if W_SHARED else offs + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + # Start the uncached residual load first, then issue the other combine + # loads before consuming any of them. + res = tl.load(res_ptr + row * stride_res + offs, mask_inner, other=0.0) + inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0) + block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0) + inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC) + inj = tl.sum(tl.where(offs_hc == stream, inj, 0.0)) + # Round the materialized combine result before normalization. This matches + # the unfused combine -> RMSNorm boundary. + out = (res.to(tl.float32) + block.to(tl.float32) * inj).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner) + + out = out.to(tl.float32) + # Keep the two-axis reduction: flattening the padded tile is ~40% slower + # at decode sizes. + sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0) + rrms = tl.rsqrt(sum_sq / HC_DIM + EPS) + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + + # Loading the weight earlier helps decode but keeps the tile live across + # the reduction and regresses larger batches, so defer it to the norm. + w = tl.load(w_ptr + w_offs, mask_inner, other=0.0) + y = out * rrms + y += y * w.to(tl.float32) + tl.store(y_ptr + row * stride_y + offs, y, mask_inner) + + +def _hc_combine_norm( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, +) -> tuple[torch.Tensor, torch.Tensor]: + N, DIM = residual.shape + assert DIM % hc_count == 0 + hc_dim = DIM // hc_count + assert block_output.shape == (N, hc_dim) + assert injection_logits.shape == (N, hc_count) + assert residual.stride(1) == 1 + assert block_output.stride(1) == 1 + assert injection_logits.stride(1) == 1 + assert norm_weight.is_contiguous() + assert norm_weight.numel() in (hc_dim, DIM) + + out = residual.new_empty(residual.shape) + y = residual.new_empty(residual.shape) + # The M=1 SM70 path is register-bound with the generic 512-wide tile. + # A 1024-wide tile keeps identical reduction/rounding results and is + # measurably faster on V100; retain the generic tile for larger batches. + BLOCK_SIZE = 1024 if N == 1 and current_platform.is_device_capability(70) else 512 + _hc_combine_norm_kernel[(N, hc_count)]( + block_output, + residual, + injection_logits, + norm_weight, + out, + y, + block_output.stride(0), + residual.stride(0), + injection_logits.stride(0), + out.stride(0), + y.stride(0), + hc_dim, + hc_count, + W_SHARED=norm_weight.numel() == hc_dim, + EPS=eps, + BLOCK_SIZE=BLOCK_SIZE, + launch_pdl=current_platform.is_arch_support_pdl(), + num_warps=4, + ) + return out, y + + +def _same_shape_fake(x: torch.Tensor, *args) -> torch.Tensor: + return x.new_empty(x.shape) + + +def _hc_gate_mix_fake( + x: torch.Tensor, gate: torch.Tensor, hc_count: int +) -> torch.Tensor: + del gate + return x.new_empty((x.shape[0], x.shape[1] // hc_count)) + + +def _hc_combine_fake( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + hc_count: int, +) -> torch.Tensor: + del block_output, injection_logits, hc_count + return residual.new_empty(residual.shape) + + +def _hc_combine_norm_fake( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, +) -> tuple[torch.Tensor, torch.Tensor]: + del block_output, injection_logits, norm_weight, eps, hc_count + return residual.new_empty(residual.shape), residual.new_empty(residual.shape) + + +direct_register_custom_op( + op_name="qwen4_exp_grouped_gemma_rmsnorm", + op_func=_grouped_gemma_rmsnorm, + fake_impl=_same_shape_fake, +) +direct_register_custom_op( + op_name="qwen4_exp_hc_silu", + op_func=_hc_silu, + fake_impl=_same_shape_fake, +) +direct_register_custom_op( + op_name="qwen4_exp_hc_gate_mix", + op_func=_hc_gate_mix, + fake_impl=_hc_gate_mix_fake, +) +direct_register_custom_op( + op_name="qwen4_exp_hc_combine", + op_func=_hc_combine, + fake_impl=_hc_combine_fake, +) +direct_register_custom_op( + op_name="qwen4_exp_hc_combine_norm", + op_func=_hc_combine_norm, + fake_impl=_hc_combine_norm_fake, +) + + +def grouped_gemma_rmsnorm( + x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int +) -> torch.Tensor: + return torch.ops.vllm.qwen4_exp_grouped_gemma_rmsnorm(x, weight, eps, num_groups) + + +def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor: + return torch.ops.vllm.qwen4_exp_hc_silu(x, hc_count) + + +def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor: + return torch.ops.vllm.qwen4_exp_hc_gate_mix(x, gate, hc_count) + + +def hc_combine( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + hc_count: int, +) -> torch.Tensor: + return torch.ops.vllm.qwen4_exp_hc_combine( + residual, block_output, injection_logits, hc_count + ) + + +def hc_combine_norm( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.vllm.qwen4_exp_hc_combine_norm( + residual, + block_output, + injection_logits, + norm_weight, + eps, + hc_count, + ) + + +__all__ = [ + "grouped_gemma_rmsnorm", + "hc_combine", + "hc_combine_norm", + "hc_gate_mix", + "hc_silu", +] diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py new file mode 100644 index 0000000000..d62a16385d --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -0,0 +1,1473 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton kernels for the Qwen4Exp weight-free QSA path.""" + +from __future__ import annotations + +import math +import os + +import torch + +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton + +logger = init_logger(__name__) + +_LOGITS_WORKSPACE_BYTES = 128 * 1024 * 1024 +_TOPK_WORKSPACE_BYTES = 1024 * 1024 +_SM70_INDEXER_CUBLAS = os.getenv("VLLM_SM70_QSA_INDEXER_CUBLAS", "1") == "1" +_SM70_INDEXER_SCORE_TILE_BYTES = ( + int(os.getenv("VLLM_SM70_QSA_INDEXER_SCORE_TILE_MB", "64")) * 1024 * 1024 +) +_SM70_INDEXER_CUBLAS_MIN_ROWS = int( + os.getenv("VLLM_SM70_QSA_INDEXER_CUBLAS_MIN_ROWS", "512") +) +_SM70_INDEXER_CUBLAS_MIN_SCORE_ELEMENTS = int( + os.getenv("VLLM_SM70_QSA_INDEXER_CUBLAS_MIN_SCORE_ELEMENTS", str(1024**2)) +) + + +@triton.jit +def _qsa_mqa_paged_kernel( + q_ptr, + k_cache_ptr, + page_table_ptr, + token_to_req_ptr, + query_positions_ptr, + sequence_lengths_ptr, + visible_blocks_ptr, + logits_ptr, + stride_q_row, + stride_q_head, + stride_q_dim, + stride_cache_block, + stride_cache_token, + stride_cache_dim, + stride_table_req, + stride_table_page, + stride_logits_row, + num_rows, + num_columns, + num_pages, + num_requests, + score_divisor, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_D: tl.constexpr, + TILES_PER_PROG: tl.constexpr, + STAGES: tl.constexpr, + MAX_N: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, +) -> None: + row = tl.program_id(0) + dims = tl.arange(0, BLOCK_D) + heads = tl.arange(0, MAX_N) + request = tl.load(token_to_req_ptr + row) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + query_position = tl.load(query_positions_ptr + row) + sequence_length = tl.load( + sequence_lengths_ptr + safe_request, + mask=(request >= 0) & (request < num_requests), + other=0, + ) + visible = tl.minimum( + (query_position + 1) // COMPRESS_RATIO, + sequence_length // COMPRESS_RATIO, + ) + if tl.program_id(1) == 0: + tl.store(visible_blocks_ptr + row, visible) + tile_start = tl.program_id(1) * TILES_PER_PROG + # Top-k is bounded by visible_blocks, so columns beyond it need no value. + if tile_start * BLOCK_N >= visible: + return + tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(visible, BLOCK_N)) + tile_end = tl.minimum(tile_end, tl.cdiv(num_columns, BLOCK_N)) + + # Pad the small head axis to a tensor-core-compatible N dimension. + query = tl.load( + q_ptr + + row * stride_q_row + + heads[None, :] * stride_q_head + + dims[:, None] * stride_q_dim, + mask=(heads[None, :] < NUM_HEADS) & (dims[:, None] < HEAD_DIM), + other=0.0, + ) + column_offsets = tl.arange(0, BLOCK_N) + for tile in tl.range(tile_start, tile_end, num_stages=STAGES): + columns = tile * BLOCK_N + column_offsets + live = columns < visible + logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1) + page_offset = columns % PAGE_SIZE + physical_page = tl.load( + page_table_ptr + + safe_request * stride_table_req + + logical_page * stride_table_page, + mask=live, + other=-1, + ) + page_valid = live & (physical_page >= 0) & (physical_page < num_pages) + # physical_page * block stride can overflow int32 for large caches. + safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64) + keys = tl.load( + k_cache_ptr + + safe_physical_page[:, None] * stride_cache_block + + page_offset[:, None] * stride_cache_token + + dims[None, :] * stride_cache_dim, + mask=page_valid[:, None] & (dims[None, :] < HEAD_DIM), + other=0.0, + eviction_policy="evict_first", + ) + scores = tl.dot(keys, query, out_dtype=tl.float32) + scores = tl.where(heads[None, :] < NUM_HEADS, tl.maximum(scores, 0.0), 0.0) + score = tl.sum(scores, axis=1) / score_divisor + tl.store( + logits_ptr + row * stride_logits_row + columns, + tl.where(page_valid, score, -float("inf")), + mask=live & (columns < num_columns), + ) + + +@triton.jit +def _qsa_visible_blocks_kernel( + token_to_req_ptr, + query_positions_ptr, + sequence_lengths_ptr, + visible_blocks_ptr, + rows, + num_requests, + COMPRESS_RATIO: tl.constexpr, +) -> None: + row = tl.program_id(0) + request = tl.load(token_to_req_ptr + row, mask=row < rows, other=-1) + valid_request = (request >= 0) & (request < num_requests) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + query_position = tl.load(query_positions_ptr + row, mask=row < rows, other=-1) + sequence_length = tl.load( + sequence_lengths_ptr + safe_request, + mask=(row < rows) & valid_request, + other=0, + ) + visible = tl.minimum( + (query_position + 1) // COMPRESS_RATIO, + sequence_length // COMPRESS_RATIO, + ) + tl.store( + visible_blocks_ptr + row, + tl.where(valid_request, tl.maximum(visible, 0), 0), + mask=row < rows, + ) + + +@triton.jit +def _qsa_gather_single_request_keys_kernel( + cache_ptr, + page_table_ptr, + keys_ptr, + valid_ptr, + stride_cache_page, + stride_cache_token, + stride_cache_dim, + stride_table_page, + stride_keys_row, + columns, + num_pages, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, +) -> None: + column = tl.program_id(0) + dims = tl.arange(0, BLOCK_D) + logical_page = column // PAGE_SIZE + page_offset = column % PAGE_SIZE + physical_page = tl.load( + page_table_ptr + logical_page * stride_table_page, + mask=(column < columns) & (logical_page < PAGE_TABLE_WIDTH), + other=-1, + ) + valid = (column < columns) & (physical_page >= 0) & (physical_page < num_pages) + safe_page = tl.maximum(physical_page, 0).to(tl.int64) + key = tl.load( + cache_ptr + + safe_page * stride_cache_page + + page_offset * stride_cache_token + + dims * stride_cache_dim, + mask=valid & (dims < HEAD_DIM), + other=0.0, + ) + tl.store( + keys_ptr + column * stride_keys_row + dims, + key, + mask=(column < columns) & (dims < HEAD_DIM), + ) + tl.store(valid_ptr + column, valid, mask=column < columns) + + +@triton.jit +def _qsa_relu_headsum_visible_kernel( + score_ptr, + visible_ptr, + key_valid_ptr, + logits_ptr, + stride_score_row, + stride_score_column, + stride_logits_row, + width, + column_start, + score_divisor, + NUM_HEADS: tl.constexpr, + BLOCK_N: tl.constexpr, +) -> None: + row = tl.program_id(0) + columns = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + absolute_columns = column_start + columns + visible = tl.load(visible_ptr + row) + key_valid = tl.load( + key_valid_ptr + absolute_columns, + mask=columns < width, + other=0, + ).to(tl.int1) + live = (columns < width) & (absolute_columns < visible) & key_valid + score = tl.zeros((BLOCK_N,), dtype=tl.float32) + for head in range(NUM_HEADS): + values = tl.load( + score_ptr + + (row * NUM_HEADS + head) * stride_score_row + + columns * stride_score_column, + mask=columns < width, + other=0.0, + ).to(tl.float32) + score += tl.maximum(values, 0.0) + tl.store( + logits_ptr + row * stride_logits_row + absolute_columns, + tl.where(live, score / score_divisor, -float("inf")), + mask=columns < width, + ) + + +@triton.jit +def _expand_qsa_indices_kernel( + block_indices_ptr, + query_positions_ptr, + sequence_lengths_ptr, + token_to_req_ptr, + output_ptr, + stride_blocks_row, + stride_blocks_column, + stride_output_row, + stride_output_column, + rows, + num_requests, + BLOCK_TOPK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + TOKEN_TOPK: tl.constexpr, + OUTPUT_WIDTH: tl.constexpr, + COLUMN_BLOCK: tl.constexpr, +) -> None: + row = tl.program_id(0) + columns = tl.program_id(1) * COLUMN_BLOCK + tl.arange(0, COLUMN_BLOCK) + query_position = tl.load(query_positions_ptr + row) + request = tl.load(token_to_req_ptr + row) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + sequence_length = tl.load( + sequence_lengths_ptr + safe_request, + mask=(request >= 0) & (request < num_requests), + other=0, + ) + complete_blocks = tl.minimum( + tl.minimum( + (query_position + 1) // COMPRESS_RATIO, + sequence_length // COMPRESS_RATIO, + ), + BLOCK_TOPK, + ) + expanded_count = complete_blocks * COMPRESS_RATIO + tail_start = ((query_position + 1) // COMPRESS_RATIO) * COMPRESS_RATIO + tail_count = (query_position + 1) - tail_start + + is_expanded = columns < expanded_count + block_rank = columns // COMPRESS_RATIO + offset = columns % COMPRESS_RATIO + safe_rank = tl.minimum(block_rank, BLOCK_TOPK - 1) + block = tl.load( + block_indices_ptr + row * stride_blocks_row + safe_rank * stride_blocks_column, + mask=(row < rows) & is_expanded, + other=-1, + ) + expanded = block * COMPRESS_RATIO + offset + tail_offset = columns - expanded_count + is_tail = ( + (columns >= expanded_count) + & (tail_offset < tail_count) + & (tail_offset < COMPRESS_RATIO - 1) + ) + token = tl.where(is_expanded, expanded, tail_start + tail_offset) + valid = ( + (row < rows) + & (columns < OUTPUT_WIDTH) + & (is_expanded | is_tail) + & (token >= 0) + & (token < sequence_length) + ) + tl.store( + output_ptr + row * stride_output_row + columns * stride_output_column, + tl.where(valid, token, -1), + mask=(row < rows) & (columns < OUTPUT_WIDTH), + ) + + +@triton.jit +def _qsa_sparse_paged_gqa_splitk_kernel( + q_ptr, + k_cache_ptr, + v_cache_ptr, + indices_ptr, + block_table_ptr, + token_to_req_ptr, + partial_output_ptr, + partial_lse_ptr, + output_ptr, + stride_q_row, + stride_q_head, + stride_k_block, + stride_k_token, + stride_k_head, + stride_v_block, + stride_v_token, + stride_v_head, + stride_indices_row, + stride_table_req, + stride_output_row, + stride_output_head, + num_rows, + num_cache_blocks, + num_requests, + TOPK: tl.constexpr, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + GROUP_SIZE: tl.constexpr, + HEAD_DIM: tl.constexpr, + NUM_QUERY_HEADS: tl.constexpr, + NUM_SPLITS: tl.constexpr, + NUM_TILES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +) -> None: + row = tl.program_id(0) + kv_head = tl.program_id(1) + split_id = tl.program_id(2) + request = tl.load(token_to_req_ptr + row) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + + head_offsets = tl.arange(0, BLOCK_M) + dim_offsets = tl.arange(0, HEAD_DIM) + column_offsets = tl.arange(0, BLOCK_N) + first_head = kv_head * GROUP_SIZE + query = tl.load( + q_ptr + + row * stride_q_row + + (first_head + head_offsets[:, None]) * stride_q_head + + dim_offsets[None, :], + mask=head_offsets[:, None] < GROUP_SIZE, + other=0.0, + ) + + max_value = tl.full((BLOCK_M,), -1.0e20, dtype=tl.float32) + normalizer = tl.zeros((BLOCK_M,), dtype=tl.float32) + accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32) + softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634 + + # Dynamic bounds avoid padded main-loop iterations for uneven splits. + split_tile_start = split_id * NUM_TILES // NUM_SPLITS + split_tile_end = (split_id + 1) * NUM_TILES // NUM_SPLITS + for tile in range(split_tile_start, split_tile_end): + columns = tile * BLOCK_N + column_offsets + logical_token = tl.load( + indices_ptr + row * stride_indices_row + columns, + mask=columns < TOPK, + other=-1, + ) + safe_token = tl.maximum(logical_token, 0) + logical_page = safe_token // PAGE_SIZE + page_offset = safe_token % PAGE_SIZE + valid = ( + (request >= 0) + & (request < num_requests) + & (logical_token >= 0) + & (logical_page < PAGE_TABLE_WIDTH) + ) + physical_page = tl.load( + block_table_ptr + + safe_request * stride_table_req + + tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1), + mask=valid, + other=-1, + ) + valid &= (physical_page >= 0) & (physical_page < num_cache_blocks) + # physical_page * block stride can overflow int32 for large caches. + safe_page = tl.maximum(physical_page, 0).to(tl.int64) + keys = tl.load( + k_cache_ptr + + safe_page[None, :] * stride_k_block + + page_offset[None, :] * stride_k_token + + kv_head * stride_k_head + + dim_offsets[:, None], + mask=valid[None, :], + other=0.0, + ) + values = tl.load( + v_cache_ptr + + safe_page[:, None] * stride_v_block + + page_offset[:, None] * stride_v_token + + kv_head * stride_v_head + + dim_offsets[None, :], + mask=valid[:, None], + other=0.0, + ) + scores = tl.dot(query, keys) + # Scaling scores avoids re-quantizing a scaled query to BF16. + scores *= softmax_scale_log2 + scores = tl.where(valid[None, :], scores, -1.0e20) + next_max = tl.maximum(max_value, tl.max(scores, axis=1)) + alpha = tl.math.exp2(max_value - next_max) + probabilities = tl.where( + valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0 + ) + accumulator = tl.dot( + probabilities.to(values.dtype), + values, + acc=accumulator * alpha[:, None], + ) + normalizer = normalizer * alpha + tl.sum(probabilities, axis=1) + max_value = next_max + + has_values = normalizer > 0 + normalized_output = tl.where( + has_values[:, None], + accumulator / tl.maximum(normalizer[:, None], 1.0e-20), + 0.0, + ) + output_mask = head_offsets[:, None] < GROUP_SIZE + if NUM_SPLITS == 1: + tl.store( + output_ptr + + row * stride_output_row + + (first_head + head_offsets[:, None]) * stride_output_head + + dim_offsets[None, :], + normalized_output, + mask=output_mask, + ) + else: + partial_lse = tl.where( + has_values, + max_value + tl.math.log2(tl.maximum(normalizer, 1.0e-20)), + -float("inf"), + ) + tl.store( + partial_output_ptr + + ( + (split_id * num_rows + row) * NUM_QUERY_HEADS + + first_head + + head_offsets[:, None] + ) + * HEAD_DIM + + dim_offsets[None, :], + normalized_output, + mask=output_mask, + ) + tl.store( + partial_lse_ptr + + (split_id * num_rows + row) * NUM_QUERY_HEADS + + first_head + + head_offsets, + partial_lse, + mask=head_offsets < GROUP_SIZE, + ) + + +@triton.jit +def _qsa_merge_splitk_kernel( + partial_output_ptr, + partial_lse_ptr, + output_ptr, + stride_output_row, + stride_output_head, + num_rows, + HEAD_DIM: tl.constexpr, + NUM_QUERY_HEADS: tl.constexpr, + NUM_SPLITS: tl.constexpr, + BLOCK_SPLITS: tl.constexpr, +) -> None: + row = tl.program_id(0) + head = tl.program_id(1) + split_offsets = tl.arange(0, BLOCK_SPLITS) + dim_offsets = tl.arange(0, HEAD_DIM) + split_mask = split_offsets < NUM_SPLITS + lse = tl.load( + partial_lse_ptr + (split_offsets * num_rows + row) * NUM_QUERY_HEADS + head, + mask=split_mask, + other=-float("inf"), + ) + lse_max = tl.max(lse, axis=0) + has_values = lse_max > -float("inf") + shifted = tl.where(split_mask & has_values, lse - lse_max, -float("inf")) + weights = tl.math.exp2(shifted) + denominator = tl.sum(weights, axis=0) + partial_output = tl.load( + partial_output_ptr + + ((split_offsets[:, None] * num_rows + row) * NUM_QUERY_HEADS + head) + * HEAD_DIM + + dim_offsets[None, :], + mask=split_mask[:, None], + other=0.0, + ) + merged = tl.sum(partial_output * weights[:, None], axis=0) + merged = tl.where(denominator > 0, merged / denominator, 0.0) + tl.store( + output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets, + merged, + ) + + +@triton.jit +def _store_qsa_rows_kernel( + cache_ptr, + slots_ptr, + rows_ptr, + stride_cache_block, + stride_cache_token, + stride_cache_dim, + stride_rows_row, + stride_rows_dim, + num_rows, + num_blocks, + PAGE_SIZE: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK_D: tl.constexpr, +) -> None: + row = tl.program_id(0) + dims = tl.arange(0, BLOCK_D) + slot = tl.load(slots_ptr + row) + valid = (row < num_rows) & (slot >= 0) & (slot < num_blocks * PAGE_SIZE) + block = tl.maximum(slot, 0) // PAGE_SIZE + token = tl.maximum(slot, 0) % PAGE_SIZE + values = tl.load( + rows_ptr + row * stride_rows_row + dims * stride_rows_dim, + mask=valid & (dims < WIDTH), + other=0, + ) + tl.store( + cache_ptr + + block * stride_cache_block + + token * stride_cache_token + + dims * stride_cache_dim, + values, + mask=valid & (dims < WIDTH), + ) + + +@triton.jit +def _compress_qsa_groups_kernel( + raw_keys_ptr, # this step's raw key rows, straight from activations + raw_positions_ptr, # this step's per-token positions + compressor_state_cache_ptr, # per-request ring of previous raw keys + rope_cache_ptr, # packed RoPE position tail of the ring + compressor_state_table_ptr, + token_to_req_ptr, + query_start_loc_ptr, + logical_positions_ptr, + compressed_slots_ptr, + pooled_ptr, + first_positions_ptr, + stride_raw_row, + stride_raw_dim, + stride_raw_positions_row, + stride_raw_positions_dim, + stride_compressor_state_block, + stride_compressor_state_token, + stride_compressor_state_dim, + stride_rope_block, + stride_rope_token, + stride_rope_dim, + stride_compressor_state_table_req, + stride_pooled_row, + stride_pooled_dim, + stride_positions_row, + stride_positions_dim, + num_rows, + num_compressor_state_blocks, + num_requests, + COMPRESSOR_STATE_SIZE: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + LOAD_ROPE_POSITIONS: tl.constexpr, +) -> None: + row = tl.program_id(0) + dims = tl.arange(0, BLOCK_D) + request = tl.load(token_to_req_ptr + row) + end_position = tl.load(logical_positions_ptr + row) + compressed_slot = tl.load(compressed_slots_ptr + row) + valid_request = (request >= 0) & (request < num_requests) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + query_row_start = tl.load( + query_start_loc_ptr + safe_request, mask=valid_request, other=0 + ) + query_row_end = tl.load( + query_start_loc_ptr + safe_request + 1, mask=valid_request, other=0 + ) + chunk_start_position = end_position - (row - query_row_start) + compressor_state_block = tl.load( + compressor_state_table_ptr + safe_request * stride_compressor_state_table_req, + mask=valid_request, + other=-1, + ) + valid_compressor_state_block = (compressor_state_block >= 0) & ( + compressor_state_block < num_compressor_state_blocks + ) + valid_row = ( + (row < num_rows) + & valid_request + & (row >= query_row_start) + & (row < query_row_end) + & (end_position >= COMPRESS_RATIO - 1) + & (compressed_slot >= 0) + ) + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) + + # A group can span the compressor-state ring (older members) and this + # step's raw rows (members at positions >= chunk_start_position). + for group_offset in tl.range(0, COMPRESS_RATIO): + position = end_position - (COMPRESS_RATIO - 1 - group_offset) + use_raw = position >= chunk_start_position + raw_row = query_row_start + position - chunk_start_position + raw_values = tl.load( + raw_keys_ptr + raw_row * stride_raw_row + dims * stride_raw_dim, + mask=valid_row + & use_raw + & (raw_row >= query_row_start) + & (raw_row < query_row_end) + & (raw_row < num_rows) + & (dims < HEAD_DIM), + other=0.0, + ).to(tl.float32) + compressor_state_values = tl.load( + compressor_state_cache_ptr + + tl.maximum(compressor_state_block, 0).to(tl.int64) + * stride_compressor_state_block + + (position % COMPRESSOR_STATE_SIZE) * stride_compressor_state_token + + dims * stride_compressor_state_dim, + mask=valid_row + & ~use_raw + & valid_compressor_state_block + & (dims < HEAD_DIM), + other=0.0, + ).to(tl.float32) + accumulator += tl.where(use_raw, raw_values, compressor_state_values) + + tl.store( + pooled_ptr + row * stride_pooled_row + dims * stride_pooled_dim, + accumulator / COMPRESS_RATIO, + mask=(row < num_rows) & (dims < HEAD_DIM), + ) + + position_dims = tl.arange(0, 4) + first_position = end_position - COMPRESS_RATIO + 1 + if LOAD_ROPE_POSITIONS: + first_from_raw = first_position >= chunk_start_position + raw_first_row = query_row_start + first_position - chunk_start_position + raw_position_values = tl.load( + raw_positions_ptr + + raw_first_row * stride_raw_positions_row + + position_dims * stride_raw_positions_dim, + mask=valid_row + & first_from_raw + & (raw_first_row >= query_row_start) + & (raw_first_row < query_row_end) + & (raw_first_row < num_rows) + & (position_dims < 3), + other=0, + ) + compressor_state_position_values = tl.load( + rope_cache_ptr + + tl.maximum(compressor_state_block, 0).to(tl.int64) * stride_rope_block + + (first_position % COMPRESSOR_STATE_SIZE) * stride_rope_token + + position_dims * stride_rope_dim, + mask=valid_row + & ~first_from_raw + & valid_compressor_state_block + & (position_dims < 3), + other=0, + ) + position_values = tl.where( + first_from_raw, + raw_position_values, + compressor_state_position_values, + ) + else: + position_values = tl.where(valid_row, first_position, 0) + tl.store( + first_positions_ptr + + row * stride_positions_row + + position_dims * stride_positions_dim, + position_values, + mask=(row < num_rows) & (position_dims < 3), + ) + + +def _validate_mqa(q: torch.Tensor) -> None: + if q.ndim != 3 or q.shape[1] <= 0 or q.shape[2] <= 0: + raise ValueError("QSA query must be [rows, heads, head_dim]") + + +def qsa_mqa_paged( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, + token_to_req: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + compress_ratio: int, + num_columns: int | None = None, + score_scale: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute QSA scores directly from a paged compressed-key cache.""" + + _validate_mqa(q) + if not q.is_cuda or not HAS_TRITON: + raise RuntimeError("paged QSA scoring requires CUDA and Triton") + if k_cache.ndim != 4 or k_cache.shape[2] != 1: + raise ValueError("QSA cache must be [pages, page_size, 1, head_dim]") + if k_cache.shape[3] != q.shape[2]: + raise ValueError("QSA query and cache dimensions must match") + if page_table.ndim != 2: + raise ValueError("QSA page table must be two-dimensional") + if q.shape[0] and (not all(k_cache.shape[:2]) or not all(page_table.shape)): + raise ValueError("QSA paged scoring cache and page table must be nonempty") + if token_to_req.shape != (q.shape[0],): + raise ValueError("QSA request mapping must match query rows") + if query_positions.shape != (q.shape[0],): + raise ValueError("QSA query positions must match query rows") + if sequence_lengths.shape != (page_table.shape[0],): + raise ValueError("QSA sequence lengths must match page-table requests") + if compress_ratio <= 0: + raise ValueError("QSA compression ratio must be positive") + score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale + if score_divisor <= 0: + raise ValueError("QSA score scale must be positive") + + capacity = page_table.shape[1] * k_cache.shape[1] + columns = capacity if num_columns is None else num_columns + if columns < 0: + raise ValueError("QSA score width must be non-negative") + logits = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device) + visible_blocks = torch.empty(q.shape[0], dtype=torch.int32, device=q.device) + if not q.shape[0] or not columns: + return logits, visible_blocks + sm70_single_token = q.shape[0] == 1 and current_platform.is_device_capability(70) + # On V100 the GB300 decode tile leaves the 128-d scorer badly + # under-occupied. A 32-column, two-warp tile preserves the selected QSA + # blocks while exposing enough independent CTAs for the single-row path. + BLOCK_N = 32 if sm70_single_token else 64 + BLOCK_D = max(16, triton.next_power_of_2(q.shape[2])) + MAX_N = max(16, triton.next_power_of_2(q.shape[1])) + # Tuned on GB300: larger row batches provide enough parallelism to reuse Q. + tiles_per_program = 1 if q.shape[0] <= 32 else 8 + _qsa_mqa_paged_kernel[ + (q.shape[0], triton.cdiv(columns, BLOCK_N * tiles_per_program)) + ]( + q, + k_cache, + page_table, + token_to_req, + query_positions, + sequence_lengths, + visible_blocks, + logits, + q.stride(0), + q.stride(1), + q.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(3), + page_table.stride(0), + page_table.stride(1), + logits.stride(0), + q.shape[0], + columns, + k_cache.shape[0], + page_table.shape[0], + float(score_divisor), + PAGE_SIZE=k_cache.shape[1], + PAGE_TABLE_WIDTH=page_table.shape[1], + NUM_HEADS=q.shape[1], + HEAD_DIM=q.shape[2], + BLOCK_N=BLOCK_N, + BLOCK_D=BLOCK_D, + TILES_PER_PROG=tiles_per_program, + STAGES=2, + MAX_N=MAX_N, + COMPRESS_RATIO=compress_ratio, + num_warps=2, + ) + return logits, visible_blocks + + +def _qsa_indexer_cublas_shape_supported( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, +) -> bool: + """Whether the exact Qwen3.8 single-request index shape can use cuBLAS.""" + + return ( + q.dtype == torch.float16 + and k_cache.dtype == torch.float16 + and q.ndim == 3 + and q.shape[1:] == (4, 128) + and k_cache.ndim == 4 + and k_cache.shape[2:] == (1, 128) + and page_table.ndim == 2 + and page_table.shape[0] == 1 + ) + + +def _use_sm70_qsa_indexer_cublas( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, +) -> bool: + return ( + _SM70_INDEXER_CUBLAS + and current_platform.is_device_capability(70) + and q.shape[0] >= _SM70_INDEXER_CUBLAS_MIN_ROWS + and _qsa_indexer_cublas_shape_supported(q, k_cache, page_table) + ) + + +def _qsa_indexer_cublas_work_supported(rows: int, columns: int) -> bool: + return rows * columns >= _SM70_INDEXER_CUBLAS_MIN_SCORE_ELEMENTS + + +def _qsa_visible_blocks( + token_to_req: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + compress_ratio: int, +) -> torch.Tensor: + rows = query_positions.shape[0] + visible = torch.empty(rows, dtype=torch.int32, device=query_positions.device) + if rows: + _qsa_visible_blocks_kernel[(rows,)]( + token_to_req, + query_positions, + sequence_lengths, + visible, + rows, + sequence_lengths.shape[0], + COMPRESS_RATIO=compress_ratio, + num_warps=1, + ) + return visible + + +def _qsa_gather_single_request_keys( + k_cache: torch.Tensor, + page_table: torch.Tensor, + columns: int, +) -> tuple[torch.Tensor, torch.Tensor]: + keys = torch.empty( + (columns, k_cache.shape[3]), dtype=k_cache.dtype, device=k_cache.device + ) + valid = torch.empty(columns, dtype=torch.uint8, device=k_cache.device) + if columns: + _qsa_gather_single_request_keys_kernel[(columns,)]( + k_cache, + page_table[0], + keys, + valid, + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(3), + page_table.stride(1), + keys.stride(0), + columns, + k_cache.shape[0], + PAGE_SIZE=k_cache.shape[1], + PAGE_TABLE_WIDTH=page_table.shape[1], + HEAD_DIM=k_cache.shape[3], + BLOCK_D=triton.next_power_of_2(k_cache.shape[3]), + num_warps=4, + ) + return keys, valid + + +def _qsa_mqa_cublas( + q: torch.Tensor, + keys: torch.Tensor, + key_valid: torch.Tensor, + visible: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Score one request with Volta Tensor Cores and a fused QSA epilogue.""" + + rows, num_heads, head_dim = q.shape + columns = keys.shape[0] + logits = torch.empty((rows, columns), dtype=torch.float32, device=q.device) + if not rows or not columns: + return logits, visible + + q2 = q.reshape(rows * num_heads, head_dim) + bytes_per_column = max(1, q2.shape[0] * torch.float32.itemsize) + tile_columns = max( + 256, + min(columns, _SM70_INDEXER_SCORE_TILE_BYTES // bytes_per_column), + ) + tile_columns = max(256, tile_columns // 256 * 256) + for column_start in range(0, columns, tile_columns): + column_stop = min(column_start + tile_columns, columns) + width = column_stop - column_start + # FP16 inputs, FP32 accumulation/output. This still selects Volta HMMA + # through cuBLAS while avoiding an FP16 round trip before top-k. + scores = torch.mm( + q2, + keys[column_start:column_stop].t(), + out_dtype=torch.float32, + ) + _qsa_relu_headsum_visible_kernel[(rows, triton.cdiv(width, 256))]( + scores, + visible, + key_valid, + logits, + scores.stride(0), + scores.stride(1), + logits.stride(0), + width, + column_start, + math.sqrt(head_dim), + NUM_HEADS=num_heads, + BLOCK_N=256, + num_warps=4, + ) + return logits, visible + + +def expand_qsa_block_indices_cuda( + block_indices: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + token_to_req: torch.Tensor, + compress_ratio: int, + token_topk: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Expand compressed blocks and compact the causal tail of the open group.""" + + if not block_indices.is_cuda or not HAS_TRITON: + raise RuntimeError("QSA CUDA expansion requires Triton") + if token_topk % compress_ratio: + raise ValueError("QSA token top-k must be divisible by compression ratio") + block_topk = token_topk // compress_ratio + output_width = token_topk + compress_ratio - 1 + if block_indices.shape != (query_positions.numel(), block_topk): + raise ValueError("QSA compressed top-k has an invalid shape") + if token_to_req.shape != query_positions.shape: + raise ValueError("QSA request mapping must match query positions") + if sequence_lengths.ndim != 1 or not sequence_lengths.shape[0]: + raise ValueError("QSA request sequence lengths must be nonempty") + if out is None: + out = torch.empty( + (block_indices.shape[0], output_width), + dtype=torch.int32, + device=block_indices.device, + ) + elif out.shape != (block_indices.shape[0], output_width): + raise ValueError("QSA expansion output has an invalid shape") + if not block_indices.shape[0]: + return out + column_block = 256 + _expand_qsa_indices_kernel[ + (block_indices.shape[0], triton.cdiv(output_width, column_block)) + ]( + block_indices, + query_positions, + sequence_lengths, + token_to_req, + out, + block_indices.stride(0), + block_indices.stride(1), + out.stride(0), + out.stride(1), + block_indices.shape[0], + sequence_lengths.shape[0], + BLOCK_TOPK=block_topk, + COMPRESS_RATIO=compress_ratio, + TOKEN_TOPK=token_topk, + OUTPUT_WIDTH=output_width, + COLUMN_BLOCK=column_block, + num_warps=4, + ) + return out + + +def qsa_select_paged_tokens( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, + token_to_req: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + token_topk: int, + compress_ratio: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Score, select, and expand QSA indices without host synchronization.""" + + rows = q.shape[0] + output_width = token_topk + compress_ratio - 1 + if out is None: + out = torch.empty((rows, output_width), dtype=torch.int32, device=q.device) + if out.shape != (rows, output_width): + raise ValueError("QSA selection output has an invalid shape") + if not rows: + return out + + capacity_columns = page_table.shape[1] * k_cache.shape[1] + score_columns = capacity_columns + block_topk = token_topk // compress_ratio + contiguous_keys: torch.Tensor | None = None + key_valid: torch.Tensor | None = None + all_visible: torch.Tensor | None = None + if _use_sm70_qsa_indexer_cublas(q, k_cache, page_table): + all_visible = _qsa_visible_blocks( + token_to_req, + query_positions, + sequence_lengths, + compress_ratio, + ) + # cuBLAS has a rectangular host-side launch shape, unlike the paged + # Triton kernel's device-side early exit. Bound it to this chunk's live + # prefix so early-context prefill does not multiply the unused 140K + # model-capacity tail. This is one scalar sync per QSA layer. + score_columns = min( + capacity_columns, + max(block_topk, int(all_visible.max().item())), + ) + if _qsa_indexer_cublas_work_supported(rows, score_columns): + logger.info_once( + "Using SM70 QSA indexer prefill cuBLAS path " + "(single-request FP16, rows=%d, score_tile_mib=%d).", + rows, + _SM70_INDEXER_SCORE_TILE_BYTES // (1024 * 1024), + ) + # Gather this request's paged MQA keys once, then reuse them across + # every bounded logits chunk below. Generic, short-work and + # multi-request batches keep the paged Triton fallback. + contiguous_keys, key_valid = _qsa_gather_single_request_keys( + k_cache, page_table, score_columns + ) + else: + score_columns = capacity_columns + all_visible = None + rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(score_columns * 4, 1)) + chunk_rows = min(rows, rows_per_chunk) + blocks_buffer = torch.empty( + (chunk_rows, block_topk), dtype=torch.int32, device=q.device + ) + topk_workspace = torch.empty( + (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device + ) + for row_start in range(0, rows, rows_per_chunk): + row_end = min(row_start + rows_per_chunk, rows) + row_slice = slice(row_start, row_end) + if contiguous_keys is not None and key_valid is not None: + assert all_visible is not None + logits, visible_blocks = _qsa_mqa_cublas( + q[row_slice], + contiguous_keys, + key_valid, + all_visible[row_slice], + ) + else: + logits, visible_blocks = qsa_mqa_paged( + q[row_slice], + k_cache, + page_table, + token_to_req[row_slice], + query_positions[row_slice], + sequence_lengths, + compress_ratio, + ) + blocks = blocks_buffer[: row_end - row_start] + use_cooperative_topk = ( + blocks.shape[0] <= 32 + and logits.stride(0) % 4 == 0 + and current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) + ) + topk_op = ( + torch.ops._C.cooperative_topk + if use_cooperative_topk + else torch.ops._C.persistent_topk + ) + topk_op( + logits, + visible_blocks, + blocks, + topk_workspace, + block_topk, + score_columns, + ) + expand_qsa_block_indices_cuda( + blocks, + query_positions[row_slice], + sequence_lengths, + token_to_req[row_slice], + compress_ratio, + token_topk, + out[row_slice], + ) + return out + + +def qsa_sparse_paged_attention( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + logical_indices: torch.Tensor, + block_table: torch.Tensor, + token_to_req: torch.Tensor, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Run sparse GQA directly over paged FP16/BF16 K/V caches.""" + + if not q.is_cuda or not HAS_TRITON: + raise RuntimeError("paged QSA sparse attention requires CUDA and Triton") + if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape: + raise ValueError("QSA sparse attention received invalid Q/K/V shapes") + if logical_indices.ndim != 2 or logical_indices.shape[0] != q.shape[0]: + raise ValueError("QSA indices must have one row per query") + if token_to_req.shape != (q.shape[0],) or block_table.ndim != 2: + raise ValueError("QSA sparse attention metadata has invalid shapes") + if not all(k_cache.shape[:3]) or not all(block_table.shape): + raise ValueError("QSA sparse attention cache and block table must be nonempty") + if logical_indices.shape[1] <= 0: + raise ValueError("QSA sparse attention requires a positive selection width") + if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]: + raise ValueError("QSA sparse attention requires valid grouped-query heads") + head_dim = q.shape[2] + assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 + assert q.dtype == k_cache.dtype == v_cache.dtype + assert q.dtype in (torch.float16, torch.bfloat16) + assert logical_indices.dtype == block_table.dtype == torch.int32 + assert token_to_req.dtype == torch.int32 + assert q.device == k_cache.device == v_cache.device + assert q.device == logical_indices.device == block_table.device + assert q.device == token_to_req.device + assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1 + assert logical_indices.stride(1) == block_table.stride(1) == 1 + assert token_to_req.stride(0) == 1 + if out is None: + out = torch.empty_like(q) + if out.shape != q.shape: + raise ValueError("QSA sparse output must match its query") + assert out.dtype == q.dtype and out.device == q.device + assert out.stride(2) == 1 + if not q.shape[0]: + return out + + group_size = q.shape[1] // k_cache.shape[2] + block_m = triton.next_power_of_2(group_size) + base_programs = q.shape[0] * k_cache.shape[2] + block_n, target_splits, partial_warps = _qsa_sparse_launch_profile( + base_programs, + block_m, + current_platform.is_device_capability(70), + ) + + if ( + q.shape[0] == 1 + and group_size == 6 + and head_dim == 256 + and current_platform.is_device_capability(70) + ): + # Exact Qwen4Exp TP4 decode shape. Two warps preserve the existing + # split/merge arithmetic and cut the partial-kernel time on V100. + partial_warps = 2 + + num_tiles = triton.cdiv(logical_indices.shape[1], block_n) + # Avoid empty splits when the selection width is smaller than the profile. + max_useful_splits = 1 << (num_tiles.bit_length() - 1) + num_splits = min(max_useful_splits, target_splits) + + # Split=1 writes output directly and compiles out all workspace accesses. + if num_splits == 1: + partial_output = out + partial_lse = out + else: + # FP32 partials preserve accuracy when merging independently normalized + # splits. + partial_output = torch.empty( + (num_splits, *q.shape), dtype=torch.float32, device=q.device + ) + partial_lse = torch.empty( + (num_splits, q.shape[0], q.shape[1]), + dtype=torch.float32, + device=q.device, + ) + + partial_grid = (q.shape[0], k_cache.shape[2], num_splits) + _qsa_sparse_paged_gqa_splitk_kernel[partial_grid]( + q, + k_cache, + v_cache, + logical_indices, + block_table, + token_to_req, + partial_output, + partial_lse, + out, + q.stride(0), + q.stride(1), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + logical_indices.stride(0), + block_table.stride(0), + out.stride(0), + out.stride(1), + q.shape[0], + k_cache.shape[0], + block_table.shape[0], + TOPK=logical_indices.shape[1], + PAGE_SIZE=k_cache.shape[1], + PAGE_TABLE_WIDTH=block_table.shape[1], + GROUP_SIZE=group_size, + HEAD_DIM=q.shape[2], + NUM_QUERY_HEADS=q.shape[1], + NUM_SPLITS=num_splits, + NUM_TILES=num_tiles, + BLOCK_M=block_m, + BLOCK_N=block_n, + num_warps=partial_warps, + num_stages=2, + ) + if num_splits == 1: + return out + + _qsa_merge_splitk_kernel[(q.shape[0], q.shape[1])]( + partial_output, + partial_lse, + out, + out.stride(0), + out.stride(1), + q.shape[0], + HEAD_DIM=q.shape[2], + NUM_QUERY_HEADS=q.shape[1], + NUM_SPLITS=num_splits, + BLOCK_SPLITS=triton.next_power_of_2(num_splits), + num_warps=2, + num_stages=1, + ) + return out + + +def _qsa_sparse_launch_profile( + base_programs: int, + block_m: int, + is_sm70: bool, +) -> tuple[int, int, int]: + """Return BLOCK_N, target splits, and warps for sparse QSA.""" + small_profile_limit = 8 if block_m <= 8 else 4 + + # Tuned on GB300 for the Qwen-Air TP1, TP2, and TP4 attention shapes. + # Narrow tiles favor decode; wide tiles improve throughput for prefill. + if base_programs <= small_profile_limit: + block_n, target_splits, partial_warps = 16, 64, 4 + elif base_programs < 32: + block_n, target_splits, partial_warps = 16, 32, 4 + elif base_programs <= 256: + block_n, target_splits, partial_warps = 64, 8, 2 + elif base_programs <= 512: + block_n, target_splits, partial_warps = 64, 4, 2 + else: + block_n, target_splits, partial_warps = 64, 1, 2 + if is_sm70 and block_n == 64: + # Two warps serialize the D=256 tensor-core work on V100. Four warps + # restore warp-level parallelism for split and non-split prefill. + partial_warps = 4 + if base_programs >= 512: + # A 32-column tile improves the exact 512-row and 8192-row Qwen4Exp + # prefill shapes without changing small-batch or non-SM70 routes. + block_n = 32 + return block_n, target_splits, partial_warps + + +def qsa_store_cache_rows( + cache: torch.Tensor, + slot_mapping: torch.Tensor, + rows: torch.Tensor, +) -> None: + """Store fixed-width rows in a QSA cache without boolean indexing.""" + + if not cache.is_cuda or not HAS_TRITON: + raise RuntimeError("QSA CUDA cache stores require Triton") + if cache.ndim != 4 or cache.shape[2] != 1: + raise ValueError("QSA cache must be [pages, page_size, 1, width]") + if not all(cache.shape): + raise ValueError("QSA cache dimensions must be nonzero") + if rows.ndim == 3: + if rows.shape[1] != 1: + raise ValueError("QSA cache rows must have one head") + rows = rows[:, 0] + if rows.shape != (slot_mapping.numel(), cache.shape[3]): + raise ValueError("QSA cache rows and slots have incompatible shapes") + if not rows.shape[0]: + return + _store_qsa_rows_kernel[(rows.shape[0],)]( + cache, + slot_mapping, + rows, + cache.stride(0), + cache.stride(1), + cache.stride(3), + rows.stride(0), + rows.stride(1), + rows.shape[0], + cache.shape[0], + PAGE_SIZE=cache.shape[1], + WIDTH=cache.shape[3], + BLOCK_D=triton.next_power_of_2(cache.shape[3]), + num_warps=4, + ) + + +def qsa_compress_groups_with_ratio( + raw_keys: torch.Tensor, # this step's raw key rows [rows, 1, head_size] + raw_positions: torch.Tensor, # this step's positions [rows, 1, 3] int64 + compressor_state_cache: torch.Tensor, + compressor_state_block_table: torch.Tensor, + token_to_req: torch.Tensor, + query_start_loc: torch.Tensor, + logical_positions: torch.Tensor, + compressed_slots: torch.Tensor, + compress_ratio: int, + rope_cache: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pool completed groups from the compressor-state ring and raw token rows.""" + + if not raw_keys.is_cuda or not HAS_TRITON: + raise RuntimeError("QSA CUDA compression requires Triton") + rows = token_to_req.numel() + if compress_ratio <= 0: + raise ValueError("QSA compression ratio must be positive") + if raw_keys.ndim != 3 or raw_keys.shape[:2] != (rows, 1): + raise ValueError("QSA raw keys must be [rows, 1, head_size]") + if raw_positions.shape != (rows, 1, 3) or raw_positions.dtype != torch.int64: + raise ValueError("QSA raw positions must be [rows, 1, 3] int64") + if logical_positions.shape != (rows,) or compressed_slots.shape != (rows,): + raise ValueError("QSA compression metadata must match token rows") + if compressor_state_cache.ndim != 4 or compressor_state_cache.shape[2] != 1: + raise ValueError("QSA compressor-state cache has an invalid shape") + if ( + # The ring is wider than one group so speculative rows cannot alias + # onto the committed keys of the group still being collected. + compressor_state_cache.shape[1] < compress_ratio + or compressor_state_cache.shape[3] != raw_keys.shape[2] + or compressor_state_cache.dtype != raw_keys.dtype + ): + raise ValueError( + "QSA compressor-state cache does not match the compression layout" + ) + if ( + compressor_state_block_table.ndim != 2 + or compressor_state_block_table.shape[1] < 1 + ): + raise ValueError( + "QSA compressor-state block table must contain one block per request" + ) + if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2: + raise ValueError("QSA query starts must contain a terminal offset") + num_requests = query_start_loc.shape[0] - 1 + if compressor_state_block_table.shape[0] < num_requests: + raise ValueError("QSA compressor-state block table has too few request rows") + if rope_cache is not None and ( + rope_cache.ndim != 4 + or rope_cache.shape[:3] != compressor_state_cache.shape[:3] + or rope_cache.shape[3] != 3 + or rope_cache.dtype != torch.int64 + ): + raise ValueError("QSA packed position view has an invalid shape or dtype") + if rows and ( + not all(compressor_state_cache.shape) + or not all(compressor_state_block_table.shape) + ): + raise ValueError("QSA compressor-state cache and block table must be nonempty") + pooled = torch.empty( + (rows, 1, raw_keys.shape[2]), + dtype=raw_keys.dtype, + device=raw_keys.device, + ) + first_positions = torch.empty((rows, 3), dtype=torch.int64, device=raw_keys.device) + if not rows: + return pooled, first_positions + if rope_cache is None: + rope_cache = compressor_state_cache + load_rope_positions = False + else: + load_rope_positions = True + _compress_qsa_groups_kernel[(rows,)]( + raw_keys, + raw_positions, + compressor_state_cache, + rope_cache, + compressor_state_block_table, + token_to_req, + query_start_loc, + logical_positions, + compressed_slots, + pooled, + first_positions, + raw_keys.stride(0), + raw_keys.stride(2), + raw_positions.stride(0), + raw_positions.stride(2), + compressor_state_cache.stride(0), + compressor_state_cache.stride(1), + compressor_state_cache.stride(3), + rope_cache.stride(0), + rope_cache.stride(1), + rope_cache.stride(3), + compressor_state_block_table.stride(0), + pooled.stride(0), + pooled.stride(2), + first_positions.stride(0), + first_positions.stride(1), + rows, + compressor_state_cache.shape[0], + num_requests, + COMPRESSOR_STATE_SIZE=compressor_state_cache.shape[1], + COMPRESS_RATIO=compress_ratio, + HEAD_DIM=raw_keys.shape[2], + LOAD_ROPE_POSITIONS=load_rope_positions, + BLOCK_D=triton.next_power_of_2(raw_keys.shape[2]), + num_warps=4, + ) + return pooled, first_positions + + +__all__ = [ + "expand_qsa_block_indices_cuda", + "qsa_compress_groups_with_ratio", + "qsa_mqa_paged", + "qsa_select_paged_tokens", + "qsa_sparse_paged_attention", + "qsa_store_cache_rows", +] diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py new file mode 100644 index 0000000000..3c1695548e --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py @@ -0,0 +1,516 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused QSA pre-indexer kernel for Qwen4Exp.""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _norm_rope( + x, + pos_t, + pos_h, + pos_w, + cos_sin_ptr, + cos_sin_stride, + norm_weight_ptr, + eps, + IS_MROPE: tl.constexpr, + MROPE_H: tl.constexpr, + MROPE_W: tl.constexpr, +): + """Apply Gemma RMSNorm and selected-axis NeoX RoPE to register rows.""" + TILE_T: tl.constexpr = x.shape[0] + TILE_H: tl.constexpr = x.shape[1] + D: tl.constexpr = x.shape[2] + ROWS: tl.constexpr = TILE_T * TILE_H + HALF: tl.constexpr = D // 2 + QUARTER: tl.constexpr = D // 4 + pairs = tl.arange(0, QUARTER) + if IS_MROPE: + # Qwen interleaves temporal, height, and width rotary pairs. Each axis + # still indexes the same position-major cos/sin table. + h_mask = ((pairs % 3) == 1) & (pairs <= 3 * MROPE_H) + w_mask = ((pairs % 3) == 2) & (pairs <= 3 * MROPE_W) + t_mask = ~(h_mask | w_mask) + base = cos_sin_ptr + pairs[None, :] + pos_rows = (pos_t, pos_h, pos_w) + axis_masks = (t_mask, h_mask, w_mask) + cos = tl.zeros((TILE_T, QUARTER), dtype=cos_sin_ptr.dtype.element_ty) + sin = tl.zeros((TILE_T, QUARTER), dtype=cos_sin_ptr.dtype.element_ty) + for axis in tl.static_range(3): + cos += tl.load( + base + pos_rows[axis][:, None] * cos_sin_stride, + mask=axis_masks[axis][None, :], + other=0, + ) + sin += tl.load( + base + pos_rows[axis][:, None] * cos_sin_stride + QUARTER, + mask=axis_masks[axis][None, :], + other=0, + ) + else: + cos = tl.load(cos_sin_ptr + pos_t[:, None] * cos_sin_stride + pairs[None, :]) + sin = tl.load( + cos_sin_ptr + pos_t[:, None] * cos_sin_stride + QUARTER + pairs[None, :] + ) + + cos = tl.reshape( + tl.broadcast_to(cos[:, None, :], (TILE_T, TILE_H, QUARTER)), + (ROWS, QUARTER), + ) + sin = tl.reshape( + tl.broadcast_to(sin[:, None, :], (TILE_T, TILE_H, QUARTER)), + (ROWS, QUARTER), + ) + x = tl.reshape(x, (ROWS, D)).to(tl.float32) + weight = tl.load(norm_weight_ptr + tl.arange(0, D)).to(tl.float32) + 1.0 + rrms = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps) + y = (x * rrms[:, None] * weight[None, :]).to(cos.dtype) + rotated, passthrough = tl.split( + tl.permute(tl.reshape(y, (ROWS, 2, HALF)), (0, 2, 1)) + ) + r0, r1 = tl.split(tl.permute(tl.reshape(rotated, (ROWS, 2, QUARTER)), (0, 2, 1))) + out0 = r0 * cos - r1 * sin + out1 = r1 * cos + r0 * sin + rotated = tl.reshape(tl.permute(tl.join(out0, out1), (0, 2, 1)), (ROWS, HALF)) + result = tl.reshape(tl.permute(tl.join(rotated, passthrough), (0, 2, 1)), (ROWS, D)) + return tl.reshape(result, (TILE_T, TILE_H, D)) + + +@triton.jit( + do_not_specialize=[ + "num_tokens", + "num_state_blocks", + "num_compressed_blocks", + "num_k_work", + ] +) +def _qsa_pre_indexer_kernel( + q_ptr, + q_stride_token, + k_ptr, + k_stride_token, + pos_ptr, + pos_stride_axis, + pos_stride_token, + cos_sin_ptr, + q_norm_weight_ptr, + k_norm_weight_ptr, + eps, + q_out_ptr, + q_out_stride_token, + q_out_stride_head, + state_cache_ptr, + state_cache_stride_block, + state_cache_stride_token, + state_slots_ptr, + state_table_ptr, + state_table_stride_req, + query_start_loc_ptr, + logical_positions_ptr, + compressed_slots_ptr, + k_work_metadata_ptr, + compressed_cache_ptr, + compressed_cache_stride_block, + compressed_cache_stride_token, + num_tokens, + num_state_blocks, + num_compressed_blocks, + num_k_work, + HQ: tl.constexpr, + D: tl.constexpr, + TILE_T_Q: tl.constexpr, + TILE_H_Q: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + STATE_SIZE: tl.constexpr, + COMP_PAGE_SIZE: tl.constexpr, + IS_2D_POSITIONS: tl.constexpr, + IS_K_MROPE: tl.constexpr, + CACHE_HAS_ROPE_POS: tl.constexpr, + CACHE_IS_FP16: tl.constexpr, + MROPE_H: tl.constexpr, + MROPE_W: tl.constexpr, +): + pid = tl.program_id(0) + # K work occupies the first programs; the remaining programs tile Q. This + # keeps both paths in one launch while leaving their register shapes + # independent. + if pid >= num_k_work: + q_pid = pid - num_k_work + num_head_tiles: tl.constexpr = tl.cdiv(HQ, TILE_H_Q) + token_tile = q_pid // num_head_tiles + head_tile = q_pid % num_head_tiles + tokens = token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q) + heads = head_tile * TILE_H_Q + tl.arange(0, TILE_H_Q) + valid_tokens = tokens < num_tokens + valid_heads = heads < HQ + dims = tl.arange(0, D) + mask = valid_tokens[:, None, None] & valid_heads[None, :, None] + x = tl.load( + q_ptr + + tokens[:, None, None] * q_stride_token + + heads[None, :, None] * D + + dims[None, None, :], + mask=mask, + other=0.0, + ) + pos_t = tl.load(pos_ptr + tokens * pos_stride_token, mask=valid_tokens, other=0) + if IS_2D_POSITIONS: + pos_h = tl.load( + pos_ptr + pos_stride_axis + tokens * pos_stride_token, + mask=valid_tokens, + other=0, + ) + pos_w = tl.load( + pos_ptr + 2 * pos_stride_axis + tokens * pos_stride_token, + mask=valid_tokens, + other=0, + ) + else: + pos_h = pos_t + pos_w = pos_t + y = _norm_rope( + x, + pos_t, + pos_h, + pos_w, + cos_sin_ptr, + D // 2, + q_norm_weight_ptr, + eps, + IS_2D_POSITIONS, + MROPE_H, + MROPE_W, + ) + tl.store( + q_out_ptr + + tokens[:, None, None] * q_out_stride_token + + heads[None, :, None] * q_out_stride_head + + dims[None, None, :], + y, + mask=mask, + ) + + if pid < num_k_work: + # One work item owns one completed compression group. Work item zero + # additionally commits the request's current raw-K suffix below. + work_metadata = tl.load(k_work_metadata_ptr + pid * 2 + tl.arange(0, 2)) + request, work_in_request = tl.split(work_metadata) + if request < 0: + return + + query_start = tl.load(query_start_loc_ptr + request) + query_end = tl.load(query_start_loc_ptr + request + 1) + query_len = query_end - query_start + chunk_end = tl.load(logical_positions_ptr + query_end - 1) + chunk_start = chunk_end - query_len + 1 + num_groups = (chunk_end + 1) // COMPRESS_RATIO - chunk_start // COMPRESS_RATIO + dims = tl.arange(0, D) + + if work_in_request < num_groups: + first_boundary = ( + (chunk_start + COMPRESS_RATIO) // COMPRESS_RATIO + ) * COMPRESS_RATIO - 1 + end_position = first_boundary + work_in_request * COMPRESS_RATIO + boundary_token = query_start + end_position - chunk_start + valid_token = ( + (boundary_token >= query_start) + & (boundary_token < query_end) + & (boundary_token < num_tokens) + ) + compressed_slot = tl.load( + compressed_slots_ptr + boundary_token, + mask=valid_token, + other=-1, + ) + valid = ( + valid_token + & (compressed_slot >= 0) + & (compressed_slot < num_compressed_blocks * COMP_PAGE_SIZE) + ) + state_block = tl.load(state_table_ptr + request * state_table_stride_req) + state_block_valid = (state_block >= 0) & (state_block < num_state_blocks) + safe_state_block = tl.maximum(state_block, 0).to(tl.int64) + group_offsets = tl.arange(0, COMPRESS_RATIO) + source_positions = end_position - (COMPRESS_RATIO - 1) + group_offsets + source_in_chunk = source_positions >= chunk_start + source_tokens = query_start + source_positions - chunk_start + source_tokens_valid = ( + (source_tokens >= query_start) + & (source_tokens < query_end) + & (source_tokens < num_tokens) + ) + current_base = ( + k_ptr + tl.maximum(source_tokens, 0)[:, None] * k_stride_token + ) + cached_base = ( + state_cache_ptr + + safe_state_block * state_cache_stride_block + + (source_positions % STATE_SIZE)[:, None] * state_cache_stride_token + ) + # Only the first completed group can cross the chunk boundary. Select + # historical rows from the ring without issuing two masked loads. + source_base = tl.where(source_in_chunk[:, None], current_base, cached_base) + # Pointer selection obscures alignment from Triton's analysis. + source_base = tl.multiple_of(source_base, (8, 8)) + source_valid = tl.where( + source_in_chunk, source_tokens_valid, state_block_valid + ) + source = tl.load( + source_base + dims[None, :], + mask=valid & source_valid[:, None], + other=0.0, + ).to(tl.float32) + # Match the unfused path's storage-dtype pooled tensor before + # RMSNorm. V100 executes this route in FP16. + pooled = tl.sum(source, axis=0) / COMPRESS_RATIO + if CACHE_IS_FP16: + pooled = pooled.to(tl.float16).to(tl.float32) + else: + pooled = pooled.to(tl.bfloat16).to(tl.float32) + + first_position = end_position - (COMPRESS_RATIO - 1) + if CACHE_HAS_ROPE_POS: + # RoPE uses the first token in the pooled group. Its exact MRoPE + # coordinates may live in this chunk or the raw-state ring. + first_in_chunk = first_position >= chunk_start + first_token = query_start + first_position - chunk_start + first_token_valid = ( + (first_token >= query_start) + & (first_token < query_end) + & (first_token < num_tokens) + ) + safe_first_token = tl.maximum(first_token, 0) + load_current_position = first_in_chunk & first_token_valid + first_pos_t = tl.load( + pos_ptr + safe_first_token * pos_stride_token, + mask=load_current_position, + other=0, + ) + if IS_2D_POSITIONS: + first_pos_h = tl.load( + pos_ptr + pos_stride_axis + safe_first_token * pos_stride_token, + mask=load_current_position, + other=0, + ) + first_pos_w = tl.load( + pos_ptr + + 2 * pos_stride_axis + + safe_first_token * pos_stride_token, + mask=load_current_position, + other=0, + ) + else: + first_pos_h = first_pos_t + first_pos_w = first_pos_t + tail = ( + state_cache_ptr + + safe_state_block * state_cache_stride_block + + (first_position % STATE_SIZE) * state_cache_stride_token + + D + ).to(tl.pointer_type(tl.int64)) + load_cached_position = ~first_in_chunk & state_block_valid + cached_pos_t = tl.load(tail, mask=load_cached_position, other=0) + cached_pos_h = tl.load(tail + 1, mask=load_cached_position, other=0) + cached_pos_w = tl.load(tail + 2, mask=load_cached_position, other=0) + pos_t = tl.where(first_in_chunk, first_pos_t.to(tl.int64), cached_pos_t) + pos_h = tl.where(first_in_chunk, first_pos_h.to(tl.int64), cached_pos_h) + pos_w = tl.where(first_in_chunk, first_pos_w.to(tl.int64), cached_pos_w) + else: + pos_t = first_position + pos_h = first_position + pos_w = first_position + y = _norm_rope( + tl.reshape(pooled, (1, 1, D)), + pos_t + tl.arange(0, 1), + pos_h + tl.arange(0, 1), + pos_w + tl.arange(0, 1), + cos_sin_ptr, + D // 2, + k_norm_weight_ptr, + eps, + IS_K_MROPE, + MROPE_H, + MROPE_W, + ) + compressed_block = (compressed_slot // COMP_PAGE_SIZE).to(tl.int64) + compressed_row = compressed_slot % COMP_PAGE_SIZE + tl.store( + compressed_cache_ptr + + compressed_block * compressed_cache_stride_block + + compressed_row * compressed_cache_stride_token + + dims, + tl.reshape(y, (D,)), + mask=valid, + ) + + if work_in_request == 0: + # This CTA may have just read historical rows from the circular buffer. + # Keep every lane past those loads before overwriting the same ring. + tl.debug_barrier() + # One CTA per request commits only the suffix retained by the ring. + num_state_rows = tl.minimum(query_len, STATE_SIZE) + for state_offset in tl.range(0, num_state_rows): + token = query_end - num_state_rows + state_offset + valid_token = ( + (token >= query_start) & (token < query_end) & (token < num_tokens) + ) + slot = tl.load(state_slots_ptr + token, mask=valid_token, other=-1) + valid_slot = ( + valid_token & (slot >= 0) & (slot < num_state_blocks * STATE_SIZE) + ) + safe_slot = tl.maximum(slot, 0) + state_row = ( + state_cache_ptr + + (safe_slot // STATE_SIZE).to(tl.int64) * state_cache_stride_block + + (safe_slot % STATE_SIZE) * state_cache_stride_token + ) + k = tl.load( + k_ptr + tl.maximum(token, 0) * k_stride_token + dims, + mask=valid_slot, + other=0.0, + ) + tl.store(state_row + dims, k, mask=valid_slot) + if CACHE_HAS_ROPE_POS: + pos_t = tl.load( + pos_ptr + tl.maximum(token, 0) * pos_stride_token, + mask=valid_slot, + other=0, + ) + if IS_2D_POSITIONS: + pos_h = tl.load( + pos_ptr + + pos_stride_axis + + tl.maximum(token, 0) * pos_stride_token, + mask=valid_slot, + other=0, + ) + pos_w = tl.load( + pos_ptr + + 2 * pos_stride_axis + + tl.maximum(token, 0) * pos_stride_token, + mask=valid_slot, + other=0, + ) + else: + pos_h = pos_t + pos_w = pos_t + tail = (state_row + D).to(tl.pointer_type(tl.int64)) + tl.store(tail, pos_t.to(tl.int64), mask=valid_slot) + tl.store(tail + 1, pos_h.to(tl.int64), mask=valid_slot) + tl.store(tail + 2, pos_w.to(tl.int64), mask=valid_slot) + + +def qsa_pre_indexer( + q: torch.Tensor, + k: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + eps: float, + q_out: torch.Tensor, + state_cache: torch.Tensor, + state_slots: torch.Tensor, + state_block_table: torch.Tensor, + query_start_loc: torch.Tensor, + logical_positions: torch.Tensor, + compressed_cache: torch.Tensor, + compressed_slots: torch.Tensor, + k_work_metadata: torch.Tensor, + *, + compress_ratio: int, + mrope_section: tuple[int, int, int] | None, + rope_pos_offset: int | None, +) -> None: + """Normalize Q, compress K, then update the circular raw state.""" + num_tokens = q.shape[0] + if num_tokens == 0: + return + num_q_heads, head_dim = q_out.shape[1:] + assert cos_sin_cache.shape[-1] * 2 == head_dim + assert q.shape == (num_tokens, num_q_heads * head_dim) + assert k.shape == (num_tokens, head_dim) + assert q.stride(-1) == 1 + assert k.stride(-1) == 1 + assert q.dtype == k.dtype == q_out.dtype + assert q.dtype == state_cache.dtype == compressed_cache.dtype + assert q.dtype in (torch.float16, torch.bfloat16) + assert q_out.stride(-1) == 1 + assert cos_sin_cache.is_contiguous() + assert state_cache.stride(-1) == 1 + assert compressed_cache.stride(-1) == 1 + assert k_work_metadata.ndim == 2 and k_work_metadata.shape[1] == 2 + is_2d_positions = positions.ndim == 2 + is_k_mrope = bool(mrope_section) + cache_has_rope_pos = rope_pos_offset is not None + assert rope_pos_offset is None or rope_pos_offset == head_dim + if is_2d_positions: + assert positions.shape == (3, num_tokens) + assert is_k_mrope + pos_stride_axis, pos_stride_token = positions.stride() + else: + assert positions.shape == (num_tokens,) + pos_stride_axis, pos_stride_token = 0, positions.stride(0) + section = mrope_section if mrope_section is not None else (0, 0, 0) + assert len(section) == 3 + + if num_tokens <= 4096: + TILE_T_Q, TILE_H_Q = 2, 2 + else: + TILE_T_Q, TILE_H_Q = 2, 4 + num_k_work = k_work_metadata.shape[0] + num_q_work = triton.cdiv(num_tokens, TILE_T_Q) * triton.cdiv(num_q_heads, TILE_H_Q) + _qsa_pre_indexer_kernel[(num_k_work + num_q_work,)]( + q, + q.stride(0), + k, + k.stride(0), + positions, + pos_stride_axis, + pos_stride_token, + cos_sin_cache, + q_norm_weight, + k_norm_weight, + eps, + q_out, + q_out.stride(0), + q_out.stride(1), + state_cache, + state_cache.stride(0), + state_cache.stride(1), + state_slots, + state_block_table, + state_block_table.stride(0), + query_start_loc, + logical_positions, + compressed_slots, + k_work_metadata, + compressed_cache, + compressed_cache.stride(0), + compressed_cache.stride(1), + num_tokens, + state_cache.shape[0], + compressed_cache.shape[0], + num_k_work, + HQ=num_q_heads, + D=head_dim, + TILE_T_Q=TILE_T_Q, + TILE_H_Q=TILE_H_Q, + COMPRESS_RATIO=compress_ratio, + STATE_SIZE=state_cache.shape[1], + COMP_PAGE_SIZE=compressed_cache.shape[1], + IS_2D_POSITIONS=is_2d_positions, + IS_K_MROPE=is_k_mrope, + CACHE_HAS_ROPE_POS=cache_has_rope_pos, + CACHE_IS_FP16=q.dtype == torch.float16, + MROPE_H=section[1], + MROPE_W=section[2], + num_warps=1, + ) + + +__all__ = ["qsa_pre_indexer"] diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py new file mode 100644 index 0000000000..b4902d6c46 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -0,0 +1,1436 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp position-learning enhancement layers.""" + +import math +from collections.abc import Iterable, Sequence + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.mamba.abstract import MambaBase +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateDtypeCalculator, + MambaStateShapeCalculator, + is_conv_state_dim_first, +) +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + create_fp8_scale_parameter, + create_fp8_weight_parameter, + is_fp8, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.model_executor.parameter import ( + ModelWeightParameter, + PerTensorScaleParameter, +) +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, +) +from vllm.triton_utils import tl, triton +from vllm.utils.mem_utils import format_gib +from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import ( + direct_register_custom_op, + get_accelerator_view_from_cpu_tensor, +) +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.attention.backends.short_conv_attn import ( + PleShortConvAttentionBackend, + PleShortConvAttentionMetadata, +) +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + +from ..common.ple import copy_ple_embedding_shard_ + +_MASK64 = (1 << 64) - 1 +_SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 +_SPLITMIX_M1 = 0xBF58476D1CE4E5B9 +_SPLITMIX_M2 = 0x94D049BB133111EB +_PLE_LAYER_PRIME = 10007 + +logger = init_logger(__name__) + + +@triton.jit +def _apply_float32_sign_bit(value, sign_bit): + """Apply an FP8 sign bit without canonicalizing negative zero.""" + + value_bits = tl.cast(value, tl.uint32, bitcast=True) + signed_bits = value_bits | (sign_bit.to(tl.uint32) << 31) + return tl.cast(signed_bits, tl.float32, bitcast=True) + + +@triton.jit +def _e4m3fn_byte_to_float(raw): + """Decode E4M3FN bytes without requiring native FP8 on SM70.""" + + raw_i32 = raw.to(tl.int32) + sign_bit = (raw_i32 >> 7) & 1 + exponent = (raw_i32 >> 3) & 0x0F + mantissa = raw_i32 & 0x07 + mantissa_f32 = mantissa.to(tl.float32) + normal = (1.0 + mantissa_f32 * 0.125) * tl.exp2(exponent.to(tl.float32) - 7.0) + subnormal = mantissa_f32 * 0.001953125 # 2**-9 + value = _apply_float32_sign_bit( + tl.where(exponent == 0, subnormal, normal), sign_bit + ) + is_nan = (exponent == 0x0F) & (mantissa == 0x07) + return tl.where(is_nan, float("nan"), value) + + +@triton.jit +def _gather_ple_fp8_from_pinned_kernel( + weight_ptr, + ids_ptr, + scale_ptr, + output_ptr, + embedding_dim: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Gather and dequantize one PLE row per program from pinned host memory.""" + + row_idx = tl.program_id(0) + local_idx = tl.load(ids_ptr + row_idx) + offsets = tl.arange(0, BLOCK_D) + mask = offsets < embedding_dim + byte_ptr = weight_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8)) + raw = tl.load( + byte_ptr + local_idx * embedding_dim + offsets, + mask=mask, + other=0, + ) + scale = tl.load(scale_ptr).to(tl.float32) + values = _e4m3fn_byte_to_float(raw) * scale + tl.store(output_ptr + row_idx * embedding_dim + offsets, values, mask=mask) + + +def _splitmix64(value: int) -> int: + value = (value + _SPLITMIX_GAMMA) & _MASK64 + value = ((value ^ (value >> 30)) * _SPLITMIX_M1) & _MASK64 + value = ((value ^ (value >> 27)) * _SPLITMIX_M2) & _MASK64 + return (value ^ (value >> 31)) & _MASK64 + + +def _is_prime_64(value: int) -> bool: + if value < 2: + return False + for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): + if value % prime == 0: + return value == prime + exponent = value - 1 + shifts = 0 + while exponent % 2 == 0: + exponent //= 2 + shifts += 1 + for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022): + if base % value == 0: + continue + witness = pow(base, exponent, value) + if witness in (1, value - 1): + continue + for _ in range(shifts - 1): + witness = pow(witness, 2, value) + if witness == value - 1: + break + else: + return False + return True + + +def _nth_prime_after(start: int, count: int) -> int: + prime = int(start) + for _ in range(count): + candidate = prime + 1 + if candidate <= 2: + prime = 2 + continue + if candidate % 2 == 0: + candidate += 1 + while not _is_prime_64(candidate): + candidate += 2 + prime = candidate + return prime + + +class Qwen4ExpPLEGroupedNorm(nn.Module): + def __init__( + self, + hidden_size: int, + eps: float, + group_size: int | None, + dtype: torch.dtype | None, + ) -> None: + super().__init__() + if group_size is not None and hidden_size % group_size: + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by " + f"group_size ({group_size})" + ) + self.eps = eps + self.group_size = group_size + self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=dtype)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.float() + if self.group_size is None: + variance = hidden_states.square().mean(dim=-1, keepdim=True) + normalized = hidden_states * torch.rsqrt(variance + self.eps) + else: + grouped = hidden_states.unflatten( + -1, (hidden_states.shape[-1] // self.group_size, self.group_size) + ) + variance = grouped.square().mean(dim=-1, keepdim=True) + normalized = (grouped * torch.rsqrt(variance + self.eps)).flatten(-2) + return (normalized * (1.0 + self.weight.float())).to(input_dtype) + + +class Qwen4ExpPLEFp8EmbeddingMethod(QuantizeMethodBase): + """FP8 PLE embedding with one global checkpoint scale.""" + + def create_weights( + self, + layer: nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del input_size, output_size + weight_loader = extra_weight_attrs.get("weight_loader") + weight = create_fp8_weight_parameter( + sum(output_partition_sizes), input_size_per_partition, weight_loader + ) + layer.register_parameter("weight", weight) + + weight_scale = create_fp8_scale_parameter( + PerTensorScaleParameter, + output_partition_sizes, + input_size_per_partition, + None, + weight_loader, + # Keep graph inputs in the requested model dtype. In particular, + # an otherwise-FP16 graph cannot retain a BF16 scale parameter on + # SM70 because Inductor rejects BF16 graph inputs there. + scale_dtype=params_dtype, + ) + layer.register_parameter("weight_scale", weight_scale) + + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + raise NotImplementedError("PLE FP8 weights only support embedding lookup") + + def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: + pinned_lookup = getattr(layer, "embedding_lookup", None) + if pinned_lookup is not None: + return pinned_lookup(input_) + get_accelerator_weight = getattr(layer, "get_accelerator_weight", None) + weight = ( + get_accelerator_weight(input_.device) + if get_accelerator_weight is not None + else layer.weight + ) + return F.embedding(input_, weight) + + def process_weights_after_loading(self, layer: nn.Module) -> None: + prepare_accelerator_weight = getattr(layer, "prepare_accelerator_weight", None) + if prepare_accelerator_weight is not None: + prepare_accelerator_weight() + + +def _get_ple_embedding_quant_method( + quant_config: QuantizationConfig | None, + prefix: str, + *, + force_fp8_storage: bool = False, +) -> QuantizeMethodBase | None: + """Select global-scale FP8 only for quantized PLE checkpoint shards.""" + + if force_fp8_storage: + return Qwen4ExpPLEFp8EmbeddingMethod() + if not isinstance(quant_config, Fp8Config): + return None + if not quant_config.is_checkpoint_fp8_serialized: + return None + + ignored_layers = quant_config.ignored_layers + if is_layer_skipped( + prefix, + ignored_layers, + quant_config.packed_modules_mapping, + ): + return None + # PLE checkpoint shards form one runtime embedding parameter. + shard_prefix = f"{prefix}.shard_" + if any(name.startswith(shard_prefix) for name in ignored_layers): + return None + return Qwen4ExpPLEFp8EmbeddingMethod() + + +def _should_use_pinned_host_ple(config: Qwen4ExpTextConfig) -> bool: + explicit = getattr(config, "ple_offload_embedding", None) + if explicit is not None: + return bool(explicit) + capability = current_platform.get_device_capability() + return capability is not None and capability.to_int() == 70 + + +class Qwen4ExpPinnedHostEmbedding(VocabParallelEmbedding): + """TP-sharded FP8 PLE table backed directly by pinned host memory. + + The base embedding is constructed on the meta device, so the full shard is + never allocated on a GPU. Checkpoint shards copy directly into the pinned + CPU parameter. A stable UVA view is created after loading and used only for + embedding gathers. + """ + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + params_dtype: torch.dtype | None, + padding_size: int, + prefix: str, + quant_method: QuantizeMethodBase, + ) -> None: + if not is_pin_memory_available(): + raise RuntimeError("Qwen4Exp PLE host offload requires pinned host memory") + if not isinstance(quant_method, Qwen4ExpPLEFp8EmbeddingMethod): + raise NotImplementedError( + "Qwen4Exp pinned-host PLE currently requires FP8 checkpoint storage" + ) + + with torch.device("meta"): + super().__init__( + num_embeddings, + embedding_dim, + params_dtype=params_dtype, + padding_size=padding_size, + prefix=prefix, + quant_method=quant_method, + ) + + meta_weight = self._parameters.get("weight") + if not isinstance(meta_weight, torch.Tensor): + raise RuntimeError("Qwen4Exp PLE meta weight was not initialized") + host_weight = ModelWeightParameter( + data=torch.empty( + tuple(meta_weight.shape), + dtype=meta_weight.dtype, + device="cpu", + pin_memory=True, + ), + input_dim=1, + output_dim=0, + weight_loader=self.weight_loader, + ) + host_weight._vllm_keep_on_cpu = True + self.weight = host_weight + self.weight_scale = create_fp8_scale_parameter( + PerTensorScaleParameter, + [self.num_embeddings_per_partition], + self.embedding_dim, + None, + self.weight_loader, + scale_dtype=params_dtype, + ) + self._accelerator_weight_views: dict[int, torch.Tensor] = {} + self._accelerator_weight_ptrs: dict[int, int] = {} + self._output_dtype = self.weight_scale.dtype + logger.info( + "Qwen4Exp PLE shard allocated in pinned host memory: %s", + format_gib(self.weight.numel() * self.weight.element_size()), + ) + + def get_accelerator_weight(self, device: torch.device) -> torch.Tensor: + if device.type != "cuda": + raise RuntimeError( + f"Qwen4Exp pinned-host PLE requires a CUDA input, got {device}" + ) + device_index = ( + torch.accelerator.current_device_index() + if device.index is None + else device.index + ) + view = self._accelerator_weight_views.get(device_index) + if view is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Qwen4Exp PLE UVA view must be prepared before CUDA graph capture" + ) + with torch.accelerator.device_index(device_index): + view = get_accelerator_view_from_cpu_tensor(self.weight) + self._accelerator_weight_views[device_index] = view + self._accelerator_weight_ptrs[device_index] = view.data_ptr() + return view + + def prepare_accelerator_weight(self) -> None: + self.get_accelerator_weight( + torch.device("cuda", torch.accelerator.current_device_index()) + ) + + def embedding_lookup(self, input_: torch.Tensor) -> torch.Tensor: + """Gather FP8 UVA rows and emit scaled model-dtype values.""" + + device_index = ( + torch.accelerator.current_device_index() + if input_.device.index is None + else input_.device.index + ) + weight_ptr = self._accelerator_weight_ptrs.get(device_index) + if weight_ptr is None: + self.get_accelerator_weight(input_.device) + weight_ptr = self._accelerator_weight_ptrs[device_index] + output = torch.empty( + (*input_.shape, self.embedding_dim), + dtype=self._output_dtype, + device=input_.device, + ) + torch.ops.vllm.qwen4_exp_ple_pinned_gather( + input_.reshape(-1), + output.reshape(-1, self.embedding_dim), + self.weight_scale, + weight_ptr, + self.embedding_dim, + ) + return output + + +class Qwen4ExpNGramEmbedding(nn.Module): + def __init__( + self, + config: Qwen4ExpTextConfig, + embedding_dim: int, + ple_dense_layer_id: int, + max_total_tokens: int, + max_num_reqs: int, + prefix: str, + quant_config: QuantizationConfig | None = None, + params_dtype: torch.dtype | None = None, + ) -> None: + super().__init__() + self.embedding_dim = embedding_dim + self.ngram_size = int(config.ngram_size) + self.heads_per_ngram = int(config.heads_per_ngram) + self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram + if self.ngram_size < 2: + raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}") + if self.heads_per_ngram <= 0: + raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}") + if embedding_dim % self.ngram_heads: + raise ValueError( + "ple_embed_dim must be divisible by total ngram heads: " + f"{embedding_dim} % {self.ngram_heads} != 0" + ) + self.head_dim = embedding_dim // self.ngram_heads + self.eos_token_id = int(config.eos_token_id) + self.unigram_vocab_size = int(config.vocab_size) + self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512)) + if self.split_ngram_parts <= 0: + raise ValueError("split_ngram_parts must be positive") + + max_multiplier = ((1 << 63) - 1) // self.unigram_vocab_size + half_bound = max(1, max_multiplier // 2) + seed = int(getattr(config, "seed", None) or 1234) + base_seed = seed + _PLE_LAYER_PRIME * ple_dense_layer_id + multipliers = [] + for index in range(self.ngram_size): + value = base_seed + _SPLITMIX_GAMMA * (index + 1) + multipliers.append(2 * (_splitmix64(value) % half_bound) + 1) + self.register_buffer( + "layer_multipliers", + torch.tensor(multipliers, dtype=torch.long), + persistent=True, + ) + + ngram_vocab_size_base = int(config.ngram_vocab_size_base) + sizes: list[int] = [] + offsets: list[int] = [] + offset = 0 + for local_head in range(self.ngram_heads): + global_head = ple_dense_layer_id * self.ngram_heads + local_head + size = _nth_prime_after(ngram_vocab_size_base - 1, global_head + 1) + sizes.append(size) + offsets.append(offset) + offset += size + self.register_buffer( + "ngram_heads_vocab_sizes", + torch.tensor(sizes, dtype=torch.long), + persistent=True, + ) + self.register_buffer( + "ngram_heads_offsets", + torch.tensor(offsets, dtype=torch.long), + persistent=True, + ) + divisor = int(config.make_ngram_vocab_size_divisible_by) + padded_vocab_size = ((offset + divisor - 1) // divisor) * divisor + embedding_prefix = f"{prefix}.ngram_embedding" + ple_storage_dtype = str( + getattr(config, "ple_embedding_dtype", "") + ).removeprefix("torch.") + quant_method = _get_ple_embedding_quant_method( + quant_config, + embedding_prefix, + force_fp8_storage=ple_storage_dtype == "float8_e4m3fn", + ) + if _should_use_pinned_host_ple(config): + if quant_method is None: + raise NotImplementedError( + "Qwen4Exp pinned-host PLE requires FP8 checkpoint storage" + ) + self.ngram_embedding = Qwen4ExpPinnedHostEmbedding( + padded_vocab_size, + self.head_dim, + params_dtype=params_dtype, + padding_size=divisor, + prefix=embedding_prefix, + quant_method=quant_method, + ) + else: + self.ngram_embedding = VocabParallelEmbedding( + padded_vocab_size, + self.head_dim, + params_dtype=params_dtype, + padding_size=divisor, + prefix=embedding_prefix, + quant_method=quant_method, + ) + self.register_buffer( + "positions_buffer", + torch.arange(max_total_tokens, dtype=torch.int64), + persistent=False, + ) + self.register_buffer( + "padded_buffer", + torch.full( + (max_num_reqs, max_total_tokens), + self.eos_token_id, + dtype=torch.int64, + ), + persistent=False, + ) + + @staticmethod + def _shift_precompute( + tokens: torch.Tensor, eos_token_id: int + ) -> tuple[torch.Tensor, torch.Tensor]: + if tokens.dim() != 2: + raise ValueError("tokens must be a 2D tensor") + batch_size, seq_len = tokens.shape + positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64) + eos_positions = torch.where(tokens == eos_token_id, positions, -1) + previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values + previous_eos = torch.cat( + [ + eos_positions.new_full((batch_size, 1), -1), + previous_eos_inclusive[:, :-1], + ], + dim=1, + ) + return positions, positions.unsqueeze(0) - previous_eos - 1 + + @staticmethod + def _shift_apply( + tokens: torch.Tensor, + positions: torch.Tensor, + position_in_segment: torch.Tensor, + shift: int, + eos_token_id: int, + ) -> torch.Tensor: + if shift == 0: + return tokens + source = positions - shift + gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1) + shifted = tokens.gather(1, gather_indices) + valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift) + return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) + + def forward( + self, + input_ids: torch.Tensor, + query_start_loc: torch.Tensor, + ngram_context: torch.Tensor, + ) -> torch.Tensor: + input_ids = input_ids.reshape(-1).long() + query_start_loc = query_start_loc.long() + num_reqs = ngram_context.shape[0] + num_tokens = input_ids.shape[0] + # The scheduler already enforces max_num_batched_tokens and max_num_seqs. + # Python comparisons here would specialize these symbolic dimensions and + # violate V2's dynamic-shape contract during AOT compilation. + + positions = self.positions_buffer[:num_tokens] + packed = self.padded_buffer[:num_reqs] + packed.fill_(self.eos_token_id) + request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 + request_indices.clamp_(max=num_reqs - 1) + columns = (positions - query_start_loc[request_indices]).clamp( + 0, packed.shape[1] - 1 + ) + packed[request_indices, columns] = input_ids + ngram_context = ngram_context.to(device=input_ids.device, dtype=torch.long) + + context = torch.cat([ngram_context, packed], dim=-1) + positions_2d, position_in_segment = self._shift_precompute( + context, self.eos_token_id + ) + shifted = [context] + for shift in range(1, self.ngram_size): + shifted.append( + self._shift_apply( + context, + positions_2d, + position_in_segment, + shift, + self.eos_token_id, + ) + ) + adjusted_columns = columns + self.ngram_size - 1 + id_blocks = [] + for ngram in range(2, self.ngram_size + 1): + start = (ngram - 2) * self.heads_per_ngram + end = start + self.heads_per_ngram + mixed = shifted[0] * self.layer_multipliers[0] + for index in range(1, ngram): + mixed = torch.bitwise_xor( + mixed, shifted[index] * self.layer_multipliers[index] + ) + sizes = self.ngram_heads_vocab_sizes[start:end] + offsets = self.ngram_heads_offsets[start:end] + ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets + id_blocks.append(ids[request_indices, adjusted_columns]) + ngram_ids = torch.cat(id_blocks, dim=-1) + return self.ngram_embedding(ngram_ids).flatten(-2) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load hash buffers and checkpoint-split embedding rows.""" + + persistent_buffers = { + "layer_multipliers": self.layer_multipliers, + "ngram_heads_offsets": self.ngram_heads_offsets, + "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes, + } + loaded: set[str] = set() + regular_weights: list[tuple[str, torch.Tensor]] = [] + shard_prefix = "ngram_embedding.shard_" + + for name, loaded_weight in weights: + leaf_name = name.rsplit(".", 1)[-1] + if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup": + continue + if name in persistent_buffers: + buffer = persistent_buffers[name] + if buffer.shape != loaded_weight.shape: + raise ValueError( + f"Shape mismatch for {name}: expected " + f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}" + ) + buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype)) + loaded.add(name) + continue + if name.startswith(shard_prefix) and name.endswith(".weight"): + shard_text = name[len(shard_prefix) : -len(".weight")] + if not shard_text.isdigit(): + regular_weights.append((name, loaded_weight)) + continue + shard_index = int(shard_text) + if shard_index >= self.split_ngram_parts: + raise ValueError( + f"PLE embedding shard index {shard_index} exceeds " + f"split_ngram_parts={self.split_ngram_parts}" + ) + embedding = self.ngram_embedding + shard_size = ( + embedding.org_vocab_size + self.split_ngram_parts - 1 + ) // self.split_ngram_parts + checkpoint_start = shard_index * shard_size + expected_rows = max( + 0, + min(shard_size, embedding.org_vocab_size - checkpoint_start), + ) + expected_shape = (expected_rows, embedding.embedding_dim) + if tuple(loaded_weight.shape) != expected_shape: + raise ValueError( + f"Shape mismatch for PLE embedding shard {shard_index}: " + f"expected {expected_shape}, got " + f"{tuple(loaded_weight.shape)}" + ) + copy_ple_embedding_shard_( + embedding.weight.data, + loaded_weight, + checkpoint_start=checkpoint_start, + tp_start=embedding.shard_indices.org_vocab_start_index, + tp_end=embedding.shard_indices.org_vocab_end_index, + ) + loaded.add("ngram_embedding.weight") + continue + regular_weights.append((name, loaded_weight)) + + if regular_weights: + loaded.update(AutoWeightsLoader(self).load_weights(regular_weights)) + return loaded + + +class Qwen4ExpPLELayer(nn.Module, MambaBase): + def __init__( + self, + config: Qwen4ExpTextConfig, + vllm_config: VllmConfig, + layer_idx: int = 0, + ple_dense_layer_id: int | None = None, + prefix: str = "", + ) -> None: + super().__init__() + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.model_config: ModelConfig = model_config + self.cache_config: CacheConfig = cache_config + self.layer_idx = layer_idx + self.ple_dense_layer_id = ( + int(ple_dense_layer_id) + if ple_dense_layer_id is not None + else int(layer_idx) + ) + self.prefix = prefix + self.hidden_size = int(config.hidden_size) + self.hc_count = config.hc_count + self.hc_hidden_size = self.hidden_size * self.hc_count + self.conv_kernel_size = int(config.ple_conv_kernel_size) + self.short_conv_dilation = int(config.ngram_size) + self.conv_state_len = (self.conv_kernel_size - 1) * self.short_conv_dilation + self.num_spec_tokens = vllm_config.num_speculative_tokens + self.activation = "silu" + self.ple_embedding: nn.Module = Qwen4ExpNGramEmbedding( + config, + int(config.ple_embed_dim), + self.ple_dense_layer_id, + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.scheduler_config.max_num_seqs, + f"{prefix}.ple_embedding", + quant_config=quant_config, + params_dtype=model_config.dtype, + ) + self.key_proj = ReplicatedLinear( + int(config.ple_embed_dim), + self.hc_hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.key_proj", + ) + self.value_proj = ReplicatedLinear( + int(config.ple_embed_dim), + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.value_proj", + ) + norm_args = ( + self.hc_hidden_size, + config.rms_norm_eps, + self.hidden_size, + model_config.dtype, + ) + self.norm_key = Qwen4ExpPLEGroupedNorm(*norm_args) + self.norm_query = Qwen4ExpPLEGroupedNorm(*norm_args) + self.norm_conv = Qwen4ExpPLEGroupedNorm(*norm_args) + self.conv1d = nn.Conv1d( + self.hc_hidden_size, + self.hc_hidden_size, + self.conv_kernel_size, + groups=self.hc_hidden_size, + padding=self.conv_state_len, + dilation=self.short_conv_dilation, + bias=False, + dtype=model_config.dtype, + ) + nn.init.zeros_(self.conv1d.weight) + self.conv1d.weight._no_reinit = True + self.kv_cache = (torch.tensor([]),) + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def _get_embedding_weight_scale(self) -> torch.Tensor | None: + embedding = getattr(self.ple_embedding, "ngram_embedding", None) + return getattr(embedding, "weight_scale", None) + + def _dequantize_embeddings( + self, + embeddings: torch.Tensor, + output_dtype: torch.dtype, + ) -> torch.Tensor: + """Dequantize PLE lookup output.""" + + if not is_fp8(embeddings): + return embeddings + weight_scale = self._get_embedding_weight_scale() + if weight_scale is None: + raise RuntimeError("FP8 PLE embedding is missing its global scale") + if weight_scale.device != embeddings.device: + raise RuntimeError("FP8 PLE embedding scale must be on the output device") + return embeddings.to(output_dtype) * weight_scale.to(output_dtype) + + @property + def mamba_type(self) -> MambaAttentionBackendEnum: + return MambaAttentionBackendEnum.SHORT_CONV + + @property + def is_kv_cache_tp_replicated(self) -> bool: + return True + + def get_attn_backend(self) -> type[PleShortConvAttentionBackend]: + return PleShortConvAttentionBackend + + def get_state_dtype(self) -> tuple[torch.dtype, ...]: + return MambaStateDtypeCalculator.short_conv_state_dtype( + self.model_config.dtype, self.cache_config.mamba_cache_dtype + ) + + def get_state_shape(self) -> Sequence[tuple[int, ...]]: + return MambaStateShapeCalculator.short_conv_state_shape( + tp_world_size=1, + intermediate_size=self.hc_hidden_size, + conv_kernel=self.conv_state_len + 1, + num_spec=self.num_spec_tokens, + ) + + def _apply_norm( + self, norm: Qwen4ExpPLEGroupedNorm, hidden_states: torch.Tensor + ) -> torch.Tensor: + shape = hidden_states.shape + return norm(hidden_states.flatten(-2)).reshape(shape) + + def _short_conv_fallback(self, inputs: torch.Tensor) -> torch.Tensor: + # Profiling / CUDA graph capture only; conv state is not updated. + inputs_t = inputs.transpose(0, 1).unsqueeze(0) + output = self.conv1d(inputs_t)[..., : inputs_t.size(-1)] + return F.silu(output).squeeze(0).transpose(0, 1) + + def _short_conv_dilated_decode_batched( + self, + x_d: torch.Tensor, + conv_state: torch.Tensor, + conv_weights: torch.Tensor, + state_indices_tensor_d: torch.Tensor, + has_initial_states_d: torch.Tensor | None, + ) -> torch.Tensor: + state_indices = state_indices_tensor_d.to( + device=conv_state.device, dtype=torch.int64 + ) + # TODO: need double-check + # FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to + # slot 0 for a safe gather, then zero output and skip write-back. + valid_state = state_indices != NULL_BLOCK_ID + state_indices = torch.where( + valid_state, state_indices, torch.zeros_like(state_indices) + ) + if has_initial_states_d is None: + has_initial_state = valid_state + else: + if has_initial_states_d.numel() < state_indices_tensor_d.numel(): + raise ValueError( + "has_initial_states_d size mismatch: " + f"got {has_initial_states_d.numel()}, " + f"need >= {state_indices_tensor_d.numel()}." + ) + has_initial_state = has_initial_states_d[ + : state_indices_tensor_d.numel() + ].to(device=conv_state.device, dtype=torch.bool) + has_initial_state = has_initial_state & valid_state + + cached_state = conv_state.index_select(0, state_indices) + state = cached_state[..., : self.conv_state_len].to(x_d.dtype) + if self.conv_state_len > 0: + initial_state = torch.where( + has_initial_state.view(-1, 1, 1), + state, + torch.zeros_like(state), + ) + history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1) + else: + history = x_d.unsqueeze(-1) + + conv_output = F.conv1d( + history, + conv_weights.unsqueeze(1).contiguous(), + groups=history.size(1), + dilation=self.short_conv_dilation, + ).squeeze(-1) + output = F.silu(conv_output) + output = output * valid_state.view(-1, 1).to(output.dtype) + + if self.conv_state_len > 0: + next_state = history[..., -self.conv_state_len :] + # Padded rows are remapped to the reserved null slot. Preserve its + # existing value while writing the new states for valid rows. + existing_base_state = cached_state[..., : self.conv_state_len] + safe_next_state = torch.where( + valid_state.view(-1, 1, 1), + next_state.to(conv_state.dtype), + existing_base_state, + ) + cached_state[..., : self.conv_state_len] = safe_next_state + conv_state.index_copy_(0, state_indices, cached_state) + + return output + + def _short_conv_dilated_prefill_batched( + self, + x_p: torch.Tensor, + metadata: PleShortConvAttentionMetadata, + conv_state: torch.Tensor, + conv_weights: torch.Tensor, + state_indices_tensor_p: torch.Tensor, + num_prefills: int, + num_decode_tokens: int, + num_prefill_tokens: int, + ) -> torch.Tensor: + # ``non_spec_query_start_loc`` covers the non-spec (decode + prefill) + # requests and equals ``query_start_loc`` when spec-decode is inactive. + non_spec_query_start_loc = metadata.non_spec_query_start_loc + if non_spec_query_start_loc is None: + raise ValueError("query_start_loc is required for prefill short-conv") + query_start_loc_p = ( + non_spec_query_start_loc[-num_prefills - 1 :] - num_decode_tokens + ) + # The metadata builder guarantees that the prefill query offsets start + # at 0 and end at num_prefill_tokens. Avoid reading those values here, + # since doing so would force a device-to-host synchronization. + has_initial_states_p = metadata.has_initial_states_p + if has_initial_states_p is None: + raise ValueError("has_initial_states_p is required for prefill short-conv") + + output = torch.empty_like(x_p) + q_starts = query_start_loc_p.to(torch.int64) + if state_indices_tensor_p.numel() < num_prefills: + raise ValueError( + "state_indices_tensor_p size mismatch: " + f"got {state_indices_tensor_p.numel()}, " + f"need >= {num_prefills}." + ) + if has_initial_states_p.numel() < num_prefills: + raise ValueError( + "has_initial_states_p size mismatch: " + f"got {has_initial_states_p.numel()}, " + f"need >= {num_prefills}." + ) + if num_prefills == 0 or x_p.numel() == 0: + return output + lengths = q_starts[1:] - q_starts[:-1] + # Use the CPU-computed packing width from the metadata builder instead + # of synchronizing on lengths.max(). + max_len = metadata.max_prefill_query_len + if max_len <= 0: + return output + + hidden_size = x_p.shape[1] + positions = torch.arange( + num_prefill_tokens, device=x_p.device, dtype=torch.int64 + ) + req_indices = torch.searchsorted(q_starts[1:], positions, right=True) + col_indices = positions - q_starts[req_indices] + + packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size)) + packed_tokens[req_indices, col_indices] = x_p + packed_tokens = packed_tokens.transpose(1, 2).contiguous() + + state_indices = state_indices_tensor_p[:num_prefills].to( + device=conv_state.device, dtype=torch.int64 + ) + valid_state = state_indices != NULL_BLOCK_ID + state_indices = torch.where( + valid_state, state_indices, torch.zeros_like(state_indices) + ) + has_initial = has_initial_states_p[:num_prefills].to( + device=conv_state.device, dtype=torch.bool + ) + if self.conv_state_len > 0: + if conv_state.shape[0] == 0: + state = conv_state.new_zeros( + (num_prefills, hidden_size, self.conv_state_len), + dtype=x_p.dtype, + ) + else: + state = conv_state.index_select(0, state_indices)[ + ..., : self.conv_state_len + ].to(x_p.dtype) + use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1) + initial_state = torch.where( + use_initial_mask, + state, + torch.zeros_like(state), + ) + history = torch.cat((initial_state, packed_tokens), dim=-1) + else: + history = packed_tokens + + conv_output = F.conv1d( + history, + conv_weights.unsqueeze(1).contiguous(), + groups=history.size(1), + dilation=self.short_conv_dilation, + ) + conv_output = F.silu(conv_output).transpose(1, 2).contiguous() + + token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64) + valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1) + valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view( + num_prefills, 1 + ) + conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0) + output.copy_(conv_output[req_indices, col_indices]) + + if self.conv_state_len > 0 and conv_state.shape[0] > 0: + state_starts = lengths.to(device=history.device, dtype=torch.int64).view( + num_prefills, 1, 1 + ) + state_offsets = torch.arange( + self.conv_state_len, device=history.device, dtype=torch.int64 + ).view(1, 1, self.conv_state_len) + next_state = history.gather( + dim=2, + index=(state_starts + state_offsets).expand(-1, history.size(1), -1), + ) + # Write back without a host synchronization. Valid, non-empty rows + # receive their new state; padding and zero-length rows keep the + # current cache value. + existing_state = conv_state.index_select(0, state_indices) + existing_base_state = existing_state[..., : self.conv_state_len] + update_mask = valid_state & (lengths.to(device=conv_state.device) > 0) + safe_next_state = torch.where( + update_mask.view(num_prefills, 1, 1), + next_state.to(conv_state.dtype), + existing_base_state, + ) + existing_state[..., : self.conv_state_len] = safe_next_state + conv_state.index_copy_(0, state_indices, existing_state) + return output + + def _short_conv_dilated_spec_batched( + self, + x_spec: torch.Tensor, + conv_state: torch.Tensor, + conv_weights: torch.Tensor, + spec_state_indices_tensor: torch.Tensor, + spec_query_start_loc: torch.Tensor, + num_accepted_tokens: torch.Tensor, + spec_query_len: int, + ) -> torch.Tensor: + """Dilated short-conv for speculative-decode (MTP) requests. + + Each spec request feeds multiple (draft + 1) query tokens. The conv + outputs are computed causally after rolling back the previous draft + state by ``num_accepted_tokens - 1``. The current candidate inputs stay + in the extended cache for the next forward, matching + ``causal_conv1d_update``. + + ``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query + length and is a Python int, so no host synchronization is needed; this + keeps the path safe for full CUDA-graph capture/replay where the buffers + are padded at the request level. + """ + num_reqs = spec_state_indices_tensor.numel() + hidden_size = x_spec.size(-1) + # Use a fixed packing width instead of synchronizing on lengths.max(). + max_len = spec_query_len + # Full CUDA graphs can pad these buffers. Only the first num_reqs + # accepted-token counts belong to actual speculative requests. + num_accepted_tokens = num_accepted_tokens[:num_reqs] + q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64) + # Keep the number of real speculative tokens on the device. + total_real_tokens = q_starts[num_reqs] + + state_indices = spec_state_indices_tensor.to( + device=conv_state.device, dtype=torch.int64 + ) + valid_state = state_indices != NULL_BLOCK_ID + state_indices = torch.where( + valid_state, state_indices, torch.zeros_like(state_indices) + ) + positions = torch.arange( + x_spec.size(0), device=x_spec.device, dtype=torch.int64 + ) + # Route graph-padded token rows to the discarded dummy request so that + # they cannot overwrite real packed data. + req_indices = torch.searchsorted(q_starts[1:], positions, right=True) + valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs) + clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0)) + col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1) + pack_req_indices = torch.where( + valid_tokens, + clamped_req_indices, + torch.full_like(req_indices, num_reqs), + ) + pack_col_indices = torch.where( + valid_tokens, col_indices, torch.zeros_like(col_indices) + ) + + # The last request row is the dummy sink for graph padding. + packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size)) + packed[pack_req_indices, pack_col_indices] = x_spec + packed = packed.transpose(1, 2).contiguous() + + if self.conv_state_len > 0: + cached_state = conv_state.index_select(0, state_indices) + rollback_offsets = num_accepted_tokens.to( + device=conv_state.device, dtype=torch.int64 + ).sub(1) + rollback_offsets = torch.where( + valid_state, + rollback_offsets.clamp_(0, max_len - 1), + torch.zeros_like(rollback_offsets), + ) + state_offsets = torch.arange( + self.conv_state_len, device=conv_state.device, dtype=torch.int64 + ).view(1, 1, self.conv_state_len) + rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets + state = cached_state.gather( + 2, rollback_indices.expand(-1, hidden_size, -1) + ).to(x_spec.dtype) + state = torch.where( + valid_state.view(num_reqs, 1, 1), + state, + torch.zeros_like(state), + ) + # Append a zeroed dummy-row state to match the [num_reqs + 1] pack. + dummy_state = state.new_zeros((1, hidden_size, self.conv_state_len)) + state_full = torch.cat((state, dummy_state), dim=0) + history = torch.cat((state_full, packed), dim=-1) + else: + history = packed + + conv_output = F.conv1d( + history, + conv_weights.unsqueeze(1).contiguous(), + groups=history.size(1), + dilation=self.short_conv_dilation, + ) + conv_output = F.silu(conv_output).transpose(1, 2).contiguous() + + output = conv_output[pack_req_indices, pack_col_indices] + output = output * valid_tokens.view(-1, 1).to(output.dtype) + + # Keep all current candidate inputs in the extended state. On the next + # target forward, ``num_accepted_tokens - 1`` selects the rollback + # window before processing the newly scheduled tokens. + if self.conv_state_len > 0: + state_capacity = self.conv_state_len + max_len - 1 + if conv_state.size(-1) < state_capacity: + raise RuntimeError( + "PLE short-conv cache cannot retain speculative tokens: " + f"got {conv_state.size(-1)}, need {state_capacity}." + ) + candidate_state = history[:num_reqs, :, 1 : state_capacity + 1] + query_lengths = q_starts[1:] - q_starts[:-1] + state_positions = torch.arange( + state_capacity, device=history.device, dtype=torch.int64 + ).view(1, 1, state_capacity) + update_lengths = (self.conv_state_len + query_lengths - 1).view( + num_reqs, 1, 1 + ) + update_mask = valid_state.view(num_reqs, 1, 1) & ( + state_positions < update_lengths + ) + existing_state = cached_state[..., :state_capacity] + next_state = torch.where( + update_mask, + candidate_state.to(conv_state.dtype), + existing_state, + ) + cached_state[..., :state_capacity] = next_state + conv_state.index_copy_(0, state_indices, cached_state) + + return output + + def _short_conv_dilated_dispatch( + self, + inputs: torch.Tensor, + metadata: PleShortConvAttentionMetadata, + conv_state: torch.Tensor, + conv_weights: torch.Tensor, + ) -> torch.Tensor: + num_prefills = metadata.num_prefills + num_decodes = metadata.num_decodes + num_decode_tokens = metadata.num_decode_tokens + num_prefill_tokens = metadata.num_prefill_tokens + has_prefill = num_prefills > 0 + has_decode = num_decodes > 0 + has_spec = metadata.spec_sequence_masks is not None + x = inputs[: metadata.num_actual_tokens] + + # Split spec / non-spec tokens. + if has_spec: + if has_prefill or has_decode: + assert metadata.spec_token_indx is not None + assert metadata.non_spec_token_indx is not None + x_spec = x.index_select(0, metadata.spec_token_indx.long()) + x_non_spec = x.index_select(0, metadata.non_spec_token_indx.long()) + else: + x_spec = x + x_non_spec = None + else: + x_spec = None + x_non_spec = x + + spec_output = None + # 1. Run the multi-query speculative-decode part. + if has_spec: + assert metadata.spec_state_indices_tensor is not None + assert metadata.spec_query_start_loc is not None + assert metadata.num_accepted_tokens is not None + spec_output = self._short_conv_dilated_spec_batched( + x_spec=x_spec, + conv_state=conv_state, + conv_weights=conv_weights, + spec_state_indices_tensor=metadata.spec_state_indices_tensor[ + : metadata.num_spec_decodes + ], + spec_query_start_loc=metadata.spec_query_start_loc, + num_accepted_tokens=metadata.num_accepted_tokens, + spec_query_len=metadata.spec_query_len, + ) + + # 2. Run regular decode and prefill requests. + conv_out_non_spec = None + state_indices_tensor = metadata.state_indices_tensor + if x_non_spec is not None: + assert state_indices_tensor is not None + if has_prefill: + state_indices_tensor_d, state_indices_tensor_p = torch.split( + state_indices_tensor, + [num_decodes, num_prefills], + dim=0, + ) + x_d, x_p = torch.split( + x_non_spec, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) + non_spec_parts: list[torch.Tensor] = [] + if has_decode: + non_spec_parts.append( + self._short_conv_dilated_decode_batched( + x_d=x_d, + conv_state=conv_state, + conv_weights=conv_weights, + state_indices_tensor_d=state_indices_tensor_d, + has_initial_states_d=metadata.has_initial_states_d, + ) + ) + non_spec_parts.append( + self._short_conv_dilated_prefill_batched( + x_p=x_p, + metadata=metadata, + conv_state=conv_state, + conv_weights=conv_weights, + state_indices_tensor_p=state_indices_tensor_p, + num_prefills=num_prefills, + num_decode_tokens=num_decode_tokens, + num_prefill_tokens=num_prefill_tokens, + ) + ) + conv_out_non_spec = torch.vstack(non_spec_parts) + else: + conv_out_non_spec = self._short_conv_dilated_decode_batched( + x_d=x_non_spec, + conv_state=conv_state, + conv_weights=conv_weights, + state_indices_tensor_d=state_indices_tensor[: x_non_spec.size(0)], + has_initial_states_d=metadata.has_initial_states_d, + ) + + # 3. Merge both parts back into the original token order. + if has_spec and conv_out_non_spec is not None: + assert metadata.spec_token_indx is not None + assert metadata.non_spec_token_indx is not None + assert spec_output is not None + output = x.new_empty((metadata.num_actual_tokens, x.size(-1))) + output.index_copy_(0, metadata.spec_token_indx, spec_output) + output.index_copy_(0, metadata.non_spec_token_indx, conv_out_non_spec) + return output + elif has_spec: + assert spec_output is not None + return spec_output + if conv_out_non_spec is None: + return x + return conv_out_non_spec + + def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + if attn_metadata is None: + return self._short_conv_fallback(inputs) + + if not isinstance(attn_metadata, dict): + raise RuntimeError( + "PLE short-conv expects per-layer attention metadata dict " + f"during inference, got {type(attn_metadata).__name__}." + ) + + layer_attn_metadata = attn_metadata.get(self.prefix) + if layer_attn_metadata is None: + raise RuntimeError( + f"Missing short-conv metadata for layer '{self.prefix}'. " + "This would bypass conv-state updates and is not allowed." + ) + if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): + raise TypeError( + "Expected PleShortConvAttentionMetadata for layer " + f"'{self.prefix}', got " + f"{type(layer_attn_metadata).__name__}." + ) + + conv_state = self.kv_cache[0] + if not is_conv_state_dim_first(): + conv_state = conv_state.transpose(-1, -2) + conv_weights = self.conv1d.weight.squeeze(1) + + state_capacity = self.conv_state_len + self.num_spec_tokens + if state_capacity > 0: + if conv_state.size(-1) < state_capacity: + raise RuntimeError( + "PLE short-conv cache is smaller than expected for " + f"dilated convolution: got {conv_state.size(-1)}, " + f"expect at least {state_capacity}." + ) + conv_state = conv_state[..., -state_capacity:] + return self._short_conv_dilated_dispatch( + inputs, + layer_attn_metadata, + conv_state, + conv_weights.to(dtype=inputs.dtype), + ) + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + query_start_loc: torch.Tensor, + ngram_context: torch.Tensor, + ) -> torch.Tensor: + input_ids = input_ids.reshape(-1) + if input_ids.shape[0] != hidden_states.shape[0]: + raise ValueError( + "PLE expects input_ids and hidden_states to have the same " + f"token length, got {input_ids.shape[0]} and " + f"{hidden_states.shape[0]}" + ) + embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context) + embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype) + key, _ = self.key_proj(embeddings) + value, _ = self.value_proj(embeddings) + token_count = hidden_states.shape[0] + key = key.reshape(token_count, self.hc_count, self.hidden_size) + query = hidden_states.reshape(token_count, self.hc_count, self.hidden_size) + key = self._apply_norm(self.norm_key, key) + query = self._apply_norm(self.norm_query, query) + gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size) + gate = torch.sigmoid(gate.sign() * gate.abs().clamp_min(1e-6).sqrt()) + gated_value = gate * value.unsqueeze(-2) + normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2) + conv_output = torch.zeros_like(normalized) + torch.ops.vllm.qwen4_exp_ple_short_conv( + normalized, + conv_output, + self.prefix, + ) + return gated_value.flatten(-2) + conv_output + + +def qwen4_exp_ple_short_conv( + inputs: torch.Tensor, + output: torch.Tensor, + layer_name: str, +) -> None: + layer = get_forward_context().no_compile_layers[layer_name] + result = layer._short_conv(inputs) + output[: result.shape[0]].copy_(result) + + +def qwen4_exp_ple_short_conv_fake( + inputs: torch.Tensor, + output: torch.Tensor, + layer_name: str, +) -> None: + return + + +def qwen4_exp_ple_pinned_gather( + input_ids: torch.Tensor, + output: torch.Tensor, + weight_scale: torch.Tensor, + weight_ptr: int, + embedding_dim: int, +) -> None: + if input_ids.numel() == 0: + return + block_d = triton.next_power_of_2(embedding_dim) + _gather_ple_fp8_from_pinned_kernel[(input_ids.numel(),)]( + weight_ptr, + input_ids, + weight_scale, + output, + embedding_dim=embedding_dim, + BLOCK_D=block_d, + num_warps=4, + ) + + +def qwen4_exp_ple_pinned_gather_fake( + input_ids: torch.Tensor, + output: torch.Tensor, + weight_scale: torch.Tensor, + weight_ptr: int, + embedding_dim: int, +) -> None: + return + + +direct_register_custom_op( + op_name="qwen4_exp_ple_pinned_gather", + op_func=qwen4_exp_ple_pinned_gather, + mutates_args=["output"], + fake_impl=qwen4_exp_ple_pinned_gather_fake, +) + + +direct_register_custom_op( + op_name="qwen4_exp_ple_short_conv", + op_func=qwen4_exp_ple_short_conv, + mutates_args=["output"], + fake_impl=qwen4_exp_ple_short_conv_fake, +) + + +__all__ = [ + "Qwen4ExpNGramEmbedding", + "Qwen4ExpPLEGroupedNorm", + "Qwen4ExpPLELayer", +] diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py new file mode 100644 index 0000000000..25e9ab967d --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -0,0 +1,523 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA QSA owner with Triton kernels.""" + +from __future__ import annotations + +from typing import ClassVar, cast + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.config.cache import CacheDType +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention.attention import ( + set_default_quant_scales, +) +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope +from vllm.model_executor.models.qwen3_next import Qwen3NextAttention +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, +) +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + canonicalize_singleton_dim_strides, + direct_register_custom_op, + kv_cache_dtype_str_to_dtype, +) +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionType, +) +from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available +from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionBackend, + FlashAttentionImpl, + FlashAttentionMetadata, + FlashAttentionMetadataBuilder, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheSpec, + get_kv_quant_mode, +) + +from ..common.qsa_cache import QSAForwardMetadata +from .indexer_qsa import QSAIndexer + + +class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder): + """Flash metadata supporting uniform decode and target-verify graphs.""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + +class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend): + """FullAttentionSpec backend used by the merged QSA owner.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [ + torch.float16, + torch.bfloat16, + ] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "float16", + "bfloat16", + ] + + @staticmethod + def get_name() -> str: + return "QWEN4_EXP_QSA_TRITON" + + @staticmethod + def get_impl_cls() -> type[Qwen4ExpQSAFlashAttentionImpl]: + return Qwen4ExpQSAFlashAttentionImpl + + @staticmethod + def get_builder_cls() -> type[Qwen4ExpQSAMetadataBuilder]: + return Qwen4ExpQSAMetadataBuilder + + @classmethod + def is_sparse(cls) -> bool: + return True + + @classmethod + def supports_kv_connector(cls) -> bool: + return False + + +class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): + """Run paged sparse GQA with the QSA Triton kernel.""" + + supports_dcp: bool = False + supports_pcp: bool = False + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + if not is_flash_attn_varlen_func_available(): + raise NotImplementedError("Qwen4Exp QSA requires FlashAttention") + if self.dcp_world_size != 1: + raise NotImplementedError( + "Qwen4Exp QSA does not support decode context parallelism" + ) + if self.kv_cache_dtype not in ("auto", "float16", "bfloat16"): + raise NotImplementedError( + "Qwen4Exp QSA requires an FP16/BF16 main KV cache" + ) + self.supports_quant_query_input = False + + def forward_qsa( + self, + layer: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: FlashAttentionMetadata, + output: torch.Tensor, + token_to_req: torch.Tensor, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + del key, value + if output_scale is not None or output_block_scale is not None: + raise NotImplementedError("QSA does not support fused output quantization") + if self.alibi_slopes is not None or self.sinks is not None: + raise NotImplementedError("QSA does not support ALiBi or attention sinks") + if self.sliding_window != (-1, -1): + raise NotImplementedError("QSA does not support sliding-window attention") + + num_tokens = attn_metadata.num_actual_tokens + output.zero_() + if num_tokens == 0: + return output + + topk_buffer = getattr(layer, "topk_indices_buffer", None) + if topk_buffer is None: + raise RuntimeError("QSA owner did not provide its top-k buffer") + logical_indices = topk_buffer[:num_tokens] + token_to_req = token_to_req[:num_tokens] + # This tree's FlashAttention cache ABI keeps K/V on dimension 1: + # [num_blocks, 2, block_size, num_kv_heads, head_size]. + key_cache, value_cache = kv_cache.unbind(1) + key_cache = canonicalize_singleton_dim_strides(key_cache) + value_cache = canonicalize_singleton_dim_strides(value_cache) + if key_cache.dtype != query.dtype or query.dtype not in ( + torch.float16, + torch.bfloat16, + ): + raise NotImplementedError("Qwen4Exp QSA requires FP16/BF16 Q/K/V") + + from .ops.qsa import qsa_sparse_paged_attention + + qsa_sparse_paged_attention( + query[:num_tokens], + key_cache, + value_cache, + logical_indices, + attn_metadata.block_table, + token_to_req, + output[:num_tokens], + ) + return output + + +class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): + """Merged Qwen full-attention owner with a QSA index side branch.""" + + supports_dcp = False + + def __init__( + self, + *, + vllm_config: VllmConfig, + config: Qwen4ExpTextConfig, + layer_id: int, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + cache_config = vllm_config.cache_config + model_config = vllm_config.model_config + if cache_config is None: + raise ValueError("Qwen4Exp QSA requires a paged KV cache") + if model_config.dtype not in (torch.float16, torch.bfloat16): + raise NotImplementedError("Qwen4Exp QSA requires FP16 or BF16") + if cache_config.cache_dtype not in ("auto", "float16", "bfloat16"): + raise NotImplementedError( + "Qwen4Exp QSA requires an FP16/BF16 main KV cache" + ) + if getattr(quant_config, "kv_cache_scheme", None) is not None: + raise NotImplementedError("Qwen4Exp QSA does not support KV quantization") + parallel_config = vllm_config.parallel_config + if ( + parallel_config.prefill_context_parallel_size > 1 + or parallel_config.decode_context_parallel_size > 1 + ): + raise NotImplementedError( + "Qwen4Exp QSA does not support context parallelism" + ) + if not getattr(config, "is_causal", True): + raise NotImplementedError("Qwen4Exp QSA requires causal decoder attention") + + self.config = config + self.hidden_size = int(config.hidden_size) + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = int(config.num_attention_heads) + if self.total_num_heads % tp_size: + raise ValueError("QSA attention heads must be divisible by TP size") + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = int(config.num_key_value_heads) + if self.total_num_kv_heads >= tp_size: + if self.total_num_kv_heads % tp_size: + raise ValueError("QSA KV heads must be divisible by TP size") + elif tp_size % self.total_num_kv_heads: + raise ValueError("TP size must be divisible by replicated QSA KV heads") + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = int(config.head_dim or self.hidden_size // self.num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.dual_chunk_attention_config = getattr( + config, "dual_chunk_attention_config", None + ) + if self.dual_chunk_attention_config is not None: + raise NotImplementedError("Qwen4Exp QSA does not support dual-chunk RoPE") + # Qwen4Exp full-attention checkpoints always pack a sigmoid output + # gate next to Q, even when an inherited config default says otherwise. + self.attn_output_gate = True + qkv_quant_config = quant_config + if quant_config is not None and quant_config.get_name() == "modelopt_fp4": + qkv_quant_config = None + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads * (1 + self.attn_output_gate), + self.total_num_kv_heads, + bias=False, + quant_config=qkv_quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + reduce_results=reduce_results, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + self.rotary_emb = get_rope( + head_size=self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters=config.rope_parameters, + ) + self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + mm_config = model_config.multimodal_config + text_only = mm_config is None or mm_config.language_model_only + mrope_section = getattr(self.rotary_emb, "mrope_section", None) + supports_mrope = bool( + type(self.rotary_emb) is MRotaryEmbedding + and mrope_section + and len(mrope_section) == 3 + and sum(mrope_section) == self.rotary_emb.rotary_dim // 2 + and getattr(self.rotary_emb, "mrope_interleaved", False) + ) + supports_dtype = getattr(self.rotary_emb, "dtype", None) in ( + torch.float16, + torch.bfloat16, + ) + self.use_fused_qk_norm_rope_gate = ( + self.attn_output_gate + and getattr(self.rotary_emb, "is_neox_style", False) + and current_platform.is_cuda() + and supports_dtype + and (text_only or supports_mrope) + ) + + self.layer_name = f"{prefix}.attn" + self.attn_type = AttentionType.DECODER + self.kv_cache_dtype = cache_config.cache_dtype + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, model_config + ) + if self.kv_cache_torch_dtype != model_config.dtype: + raise NotImplementedError( + "Qwen4Exp QSA main cache dtype must match the model dtype" + ) + self.kv_sharing_target_layer_name = None + self.kv_cache = torch.tensor([]) + set_default_quant_scales(self, register_buffer=True) + + self.attn_backend = Qwen4ExpQSAFlashAttentionBackend + self.impl = Qwen4ExpQSAFlashAttentionImpl( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + None, + None, + self.kv_cache_dtype, + None, + AttentionType.DECODER, + None, + ) + self.indexer = QSAIndexer( + vllm_config=vllm_config, + config=config, + layer_id=layer_id, + rotary_emb=self.rotary_emb, + quant_config=quant_config, + prefix=f"{prefix}.indexer", + ) + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.register_buffer( + "topk_indices_buffer", + torch.empty( + max_tokens, + self.indexer.output_width, + dtype=torch.int32, + ), + persistent=False, + ) + + static_context = vllm_config.compilation_config.static_forward_context + if self.layer_name in static_context: + raise ValueError(f"Duplicate layer name: {self.layer_name}") + static_context[self.layer_name] = self + + def get_attn_backend(self) -> type[AttentionBackend]: + return self.attn_backend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + return FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + head_size_v=self.head_dim, + dtype=self.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + ) + + def _project_qkv_gate( + self, + qkv: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Split, normalize, and rotate Q/K using this tree's Qwen3 API.""" + q_gate, key, value = qkv.split( + [self.q_size * 2, self.kv_size, self.kv_size], dim=-1 + ) + token_shape = q_gate.shape[:-1] + q_gate = q_gate.view(*token_shape, self.num_heads, 2 * self.head_dim) + query, gate = torch.chunk(q_gate, 2, dim=-1) + query = self.q_norm(query).reshape(*token_shape, self.q_size) + key = self.k_norm( + key.view(*token_shape, self.num_kv_heads, self.head_dim) + ).reshape(*token_shape, self.kv_size) + query, key = self.rotary_emb(positions, query, key) + return query, key, value, gate.reshape(*token_shape, self.q_size) + + def _run_qsa( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + ) -> None: + metadata = get_forward_context().attn_metadata + if isinstance(metadata, list): + metadata = metadata[0] + if not isinstance(metadata, dict): + output.zero_() + return + main_metadata = cast(FlashAttentionMetadata, metadata[self.layer_name]) + if self.kv_cache.numel() == 0: + raise RuntimeError("QSA main K/V cache is not bound") + + num_tokens = main_metadata.num_actual_tokens + side_metadata = cast( + QSAForwardMetadata, + metadata[self.indexer.raw_key_cache.prefix], + ) + if side_metadata.num_actual_tokens != num_tokens: + raise RuntimeError("QSA main and side metadata token counts disagree") + selected = self.indexer( + hidden_states, + positions, + self.topk_indices_buffer[:num_tokens], + ) + if selected.shape != ( + num_tokens, + self.indexer.output_width, + ): + raise RuntimeError("QSA indexer returned an invalid selection shape") + impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl) + impl.do_kv_cache_update( + self, + key, + value, + self.kv_cache, + main_metadata.slot_mapping, + ) + impl.forward_qsa( + self, + query, + key, + value, + self.kv_cache, + main_metadata, + output, + token_to_req=side_metadata.token_to_req, + ) + + def forward( + self, + positions: torch.Tensor, + output: torch.Tensor | None, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v, gate = self._project_qkv_gate(qkv, positions) + num_tokens = hidden_states.shape[0] + query = q.view(num_tokens, self.num_heads, self.head_dim) + key = k.view(num_tokens, self.num_kv_heads, self.head_dim) + value = v.view(num_tokens, self.num_kv_heads, self.head_dim) + attn_output = torch.empty_like(query) + encoded_layer_name = _encode_layer_name(self.layer_name) + if current_platform.opaque_attention_op(): + torch.ops.vllm.qwen4_exp_qsa_with_output( + hidden_states, + positions, + query, + key, + value, + attn_output, + encoded_layer_name, + ) + else: + qwen4_exp_qsa_with_output( + hidden_states, + positions, + query, + key, + value, + attn_output, + encoded_layer_name, + ) + flat_output = attn_output.view(num_tokens, -1) + if gate is not None: + flat_output = flat_output * torch.sigmoid(gate) + projected_output, _ = self.o_proj(flat_output) + if output is not None: + output.copy_(projected_output) + return projected_output + + +def qwen4_exp_qsa_with_output( + hidden_states: torch.Tensor, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + layer_name: LayerNameType, +) -> None: + """Run the complete QSA state/update/attend transaction.""" + + layer_name = _resolve_layer_name(layer_name) + layer = get_forward_context().no_compile_layers[layer_name] + if not isinstance(layer, Qwen4ExpQSAAttention): + raise TypeError(f"{layer_name} is not a Qwen4Exp QSA owner") + layer._run_qsa( + hidden_states, + positions, + query, + key, + value, + output, + ) + + +def qwen4_exp_qsa_with_output_fake( + hidden_states: torch.Tensor, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + layer_name: LayerNameType, +) -> None: + del hidden_states, positions, query, key, value, output, layer_name + + +direct_register_custom_op( + op_name="qwen4_exp_qsa_with_output", + op_func=qwen4_exp_qsa_with_output, + mutates_args=["output"], + fake_impl=qwen4_exp_qsa_with_output_fake, +) + + +__all__ = [ + "QSAIndexer", + "Qwen4ExpQSAAttention", + "Qwen4ExpQSAFlashAttentionBackend", + "Qwen4ExpQSAFlashAttentionImpl", + "qwen4_exp_qsa_with_output", +] diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index e4d402e6ca..cc5fa3c344 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -126,13 +126,16 @@ def _get_backend_priorities( *sparse_backends, ] else: - return [ + backends = [ AttentionBackendEnum.FLASH_ATTN_MLA, AttentionBackendEnum.FLASHMLA, AttentionBackendEnum.FLASHINFER_MLA, AttentionBackendEnum.TRITON_MLA, AttentionBackendEnum.FLASHMLA_SPARSE, ] + if device_capability.major == 7 and device_capability.minor == 0: + backends.insert(0, AttentionBackendEnum.GLM5_SM70_SPARSE) + return backends else: if device_capability.major == 10: return [ @@ -305,6 +308,18 @@ def get_valid_backends( return valid_backends_priorities, invalid_reasons + @classmethod + def _get_indexer_block_alignment(cls, vllm_config: VllmConfig) -> int | None: + index_kpool = getattr( + vllm_config.model_config.hf_text_config, "index_kpool", None + ) + if not index_kpool or index_kpool <= 1: + return None + + from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES + + return index_kpool * min(PAGED_MQA_PAGE_SIZES) + @classmethod def get_attn_backend_cls( cls, diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index abec3bbe88..7e777db2f1 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -518,6 +518,17 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: if model_config.is_hybrid: cls._align_hybrid_block_size(vllm_config, backend_cls) + @classmethod + def _get_indexer_block_alignment(cls, vllm_config: "VllmConfig") -> int | None: + """Extra ``block_size`` multiple a sparse indexer needs, else ``None``. + + The CUDA kpool paged-MQA indexer virtually splits each storage block + into pool pages, so ``block_size`` must be a multiple of + ``index_kpool * min(PAGED_MQA_PAGE_SIZES)`` — implemented in the CUDA + platform override. Other platforms impose no extra constraint. + """ + return None + @classmethod def _align_hybrid_block_size( cls, @@ -560,6 +571,7 @@ def _align_hybrid_block_size( num_kv_heads=model_config.get_num_kv_heads(parallel_config), head_size=model_config.get_head_size(), dtype=kv_cache_dtype, + cache_dtype_str=cache_config.cache_dtype, kv_quant_mode=kv_quant_mode, ).page_size_bytes elif cache_config.cache_dtype.startswith("turboquant_"): @@ -656,6 +668,9 @@ def _align_hybrid_block_size( mamba_page_size, kernel_block_alignment_size * attn_page_size_1_token, ) + indexer_align = cls._get_indexer_block_alignment(vllm_config) + if indexer_align: + attn_block_size = indexer_align * cdiv(attn_block_size, indexer_align) if cache_config.block_size < attn_block_size: cache_config.block_size = attn_block_size diff --git a/vllm/tool_parsers/qwen3coder_tool_parser.py b/vllm/tool_parsers/qwen3coder_tool_parser.py index 2b21f8f6ae..a14d8b2e22 100644 --- a/vllm/tool_parsers/qwen3coder_tool_parser.py +++ b/vllm/tool_parsers/qwen3coder_tool_parser.py @@ -119,6 +119,7 @@ def _reset_streaming_state(self): # Store accumulated parameters for type conversion self.accumulated_params = {} self.streaming_request = None + self.json_tool_calls_streamed = 0 def _convert_param_value( self, param_value: str, param_name: str, param_config: dict, func_name: str @@ -130,6 +131,103 @@ def _convert_param_value( param_types = extract_types_from_schema(param_schema) return coerce_to_schema_type(param_value, param_types) + def _parameter_start_positions( + self, + text: str, + param_config: dict, + ) -> list[int]: + """Find parameter tags without treating inline code as markup. + + Qwen's wire format puts parameter tags on structural boundaries. Tool + values, especially code and JSON passed to Write/Edit, may themselves + contain strings such as ````. A raw substring scan treats + those literals as new arguments and corrupts the tool call. + + Unknown names at a bare line boundary are ignored when a schema is + available, but remain accepted after a real closing tag. This preserves + malformed-output recovery without interpreting an inline literal as a + schema field. + """ + positions: list[int] = [] + search_idx = 0 + while True: + position = text.find(self.parameter_prefix, search_idx) + if position == -1: + break + search_idx = position + len(self.parameter_prefix) + + name_end = text.find(">", search_idx) + if name_end == -1: + break + name = text[search_idx:name_end] + + line_start = text.rfind("\n", 0, position) + 1 + at_line_boundary = not text[line_start:position].strip() + after_close = text[:position].rstrip().endswith(self.parameter_end_token) + after_function_header = False + if not positions: + function_start = text.rfind(self.tool_call_prefix, 0, position) + if function_start != -1: + function_header_end = text.find(">", function_start) + after_function_header = ( + function_header_end != -1 + and not text[function_header_end + 1 : position].strip() + ) + elif not text[:position].strip(): + after_function_header = True + + is_structural = at_line_boundary or after_close or after_function_header + if not is_structural: + continue + if ( + param_config + and name not in param_config + and not after_close + and not after_function_header + ): + continue + positions.append(position) + return positions + + def _find_structural_parameter_end( + self, + text: str, + *, + value_start: int, + next_parameter_start: int | None, + container_end: int | None, + allow_text_end: bool, + ) -> int | None: + """Return a closing tag only when its suffix is structural. + + Literal ```` text is common in generated XML and source + files. It closes the current value only when followed by the next real + parameter, the function/tool boundary, or (for a complete non-streamed + body) the end of the text. Streaming deliberately waits for that one- + token lookahead instead of publishing a value that may later truncate. + """ + search_idx = value_start + limit = next_parameter_start + if limit is None: + limit = container_end + if limit is None: + limit = len(text) + + while True: + position = text.find(self.parameter_end_token, search_idx) + if position == -1 or position >= limit: + return None + suffix = position + len(self.parameter_end_token) + while suffix < len(text) and text[suffix].isspace(): + suffix += 1 + if next_parameter_start is not None and suffix == next_parameter_start: + return position + if container_end is not None and suffix == container_end: + return position + if allow_text_end and suffix == len(text): + return position + search_idx = position + len(self.parameter_end_token) + def _parse_xml_function_call(self, function_call_str: str) -> ToolCall | None: # Extract function name end_index = function_call_str.find(">") @@ -140,10 +238,29 @@ def _parse_xml_function_call(self, function_call_str: str) -> ToolCall | None: param_config = find_tool_properties(self.tools, function_name) parameters = function_call_str[end_index + 1 :] param_dict = {} - for match_text in self.tool_call_parameter_regex.findall(parameters): - idx = match_text.index(">") - param_name = match_text[:idx] - param_value = str(match_text[idx + 1 :]) + param_starts = self._parameter_start_positions(parameters, param_config) + for param_index, position in enumerate(param_starts): + name_start = position + len(self.parameter_prefix) + name_end = parameters.find(">", name_start) + if name_end == -1: + continue + param_name = parameters[name_start:name_end] + value_start = name_end + 1 + next_start = ( + param_starts[param_index + 1] + if param_index + 1 < len(param_starts) + else None + ) + value_end = self._find_structural_parameter_end( + parameters, + value_start=value_start, + next_parameter_start=next_start, + container_end=len(parameters), + allow_text_end=True, + ) + if value_end is None: + value_end = next_start if next_start is not None else len(parameters) + param_value = parameters[value_start:value_end] # Remove prefix and trailing \n if param_value.startswith("\n"): param_value = param_value[1:] @@ -160,6 +277,51 @@ def _parse_xml_function_call(self, function_call_str: str) -> ToolCall | None: ), ) + @staticmethod + def _parse_json_function_call(payload: str) -> ToolCall | None: + """Parse the alternate Qwen/OpenCode JSON body inside ````. + + Some Qwen coding prompts teach ``{"name": ..., "arguments": ...}`` + while the native model template teaches XML ````. The + wrapper token is unambiguous, so accepting the JSON body keeps both + clients interoperable without treating arbitrary assistant JSON as a + tool call. + """ + try: + raw_call = json.loads(payload.strip()) + except (TypeError, json.JSONDecodeError): + return None + if not isinstance(raw_call, dict): + return None + + name = raw_call.get("name") + arguments = raw_call.get("arguments", {}) + if not isinstance(name, str) or not name: + return None + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except json.JSONDecodeError: + return None + if not isinstance(arguments, dict): + return None + + return ToolCall( + type="function", + function=FunctionCall( + name=name, + arguments=json.dumps(arguments, ensure_ascii=False), + ), + ) + + def _get_json_function_calls(self, model_output: str) -> list[ToolCall]: + calls: list[ToolCall] = [] + for match in self.tool_call_complete_regex.finditer(model_output): + parsed = self._parse_json_function_call(match.group(1)) + if parsed is not None: + calls.append(parsed) + return calls + def _get_function_calls(self, model_output: str) -> list[str]: # Find all tool calls matched_ranges = self.tool_call_regex.findall(model_output) @@ -230,6 +392,106 @@ def _pending_tool_marker_delta_prefix( return pending_and_delta[: -len(current_prefix)] return pending_and_delta + def _finish_streaming_function(self, tool_text: str) -> None: + """Finalize one streamed XML function and its JSON argument state.""" + self.json_closed = True + + func_start = tool_text.find(self.tool_call_prefix) + len(self.tool_call_prefix) + func_content_end = tool_text.find(self.function_end_token, func_start) + if func_content_end != -1: + func_content = tool_text[func_start:func_content_end] + try: + parsed_tool = self._parse_xml_function_call(func_content) + if parsed_tool and self.current_tool_index < len( + self.prev_tool_call_arr + ): + self.prev_tool_call_arr[self.current_tool_index]["arguments"] = ( + parsed_tool.function.arguments + ) + except Exception: + logger.debug( + "Failed to parse tool call during streaming: %s", + tool_text, + exc_info=True, + ) + + if self.current_tool_index < len(self.streamed_args_for_tool): + self.streamed_args_for_tool[self.current_tool_index] += "}" + else: + logger.warning( + "streamed_args_for_tool out of sync: index=%d len=%d", + self.current_tool_index, + len(self.streamed_args_for_tool), + ) + + self.in_function = False + self.accumulated_params = {} + + def _extract_json_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + ) -> DeltaMessage | None: + """Emit complete alternate-format JSON calls as atomic deltas. + + Buffering until ```` deliberately favors correctness over + partial display: downstream agents must never observe malformed JSON + arguments and then persist that broken call into the next turn. + """ + matches = list(self.tool_call_complete_regex.finditer(current_text)) + if self.json_tool_calls_streamed >= len(matches): + return None + + deltas: list[DeltaToolCall] = [] + first_new_match = None + for match_index in range(self.json_tool_calls_streamed, len(matches)): + match = matches[match_index] + parsed = self._parse_json_function_call(match.group(1)) + if parsed is None: + continue + if first_new_match is None: + first_new_match = match + + call_index = len(self.prev_tool_call_arr) + call_id = self._generate_tool_call_id() + arguments = parsed.function.arguments + self.prev_tool_call_arr.append( + { + "name": parsed.function.name, + "arguments": arguments, + } + ) + self.streamed_args_for_tool.append(arguments) + deltas.append( + DeltaToolCall( + index=call_index, + id=call_id, + function=DeltaFunctionCall( + name=parsed.function.name, + arguments=arguments, + ), + type="function", + ) + ) + + # Complete wrappers, including malformed ones, are consumed once. + self.json_tool_calls_streamed = len(matches) + if not deltas: + return None + + self.is_tool_call_started = False + self.header_sent = False + self.in_function = False + self.json_started = True + self.json_closed = True + + content = None + if first_new_match is not None and first_new_match.start() > len(previous_text): + content = current_text[len(previous_text) : first_new_match.start()] + if not content: + content = None + return DeltaMessage(content=content, tool_calls=deltas) + def extract_tool_calls( self, model_output: str, @@ -237,6 +499,22 @@ def extract_tool_calls( ) -> ExtractedToolCallInformation: # Quick check to avoid unnecessary processing if self.tool_call_prefix not in model_output: + json_tool_calls = self._get_json_function_calls(model_output) + if json_tool_calls: + self.prev_tool_call_arr = [ + { + "name": call.function.name, + "arguments": call.function.arguments, + } + for call in json_tool_calls + ] + content_index = model_output.find(self.tool_call_start_token) + content = model_output[:content_index] + return ExtractedToolCallInformation( + tools_called=True, + tool_calls=json_tool_calls, + content=content if content else None, + ) return ExtractedToolCallInformation( tools_called=False, tool_calls=[], content=model_output ) @@ -325,6 +603,17 @@ def extract_tool_calls_streaming( # Update accumulated text self.accumulated_text = current_text + # Qwen/OpenCode may emit a JSON body inside the same unambiguous + # wrapper used by the XML format. Handle a complete JSON + # call before the XML state machine consumes the wrapper and waits for + # a len(previous_text): @@ -425,13 +711,30 @@ def extract_tool_calls_streaming( self.header_sent = True self.in_function = True + # A speculative burst can finish the entire function + # before the parser has emitted its header. Do not rely on + # a later delta to revisit the already-buffered body: the + # next delta may only contain plus EOS. Parse + # the complete body now and emit its arguments together + # with the header. + complete_arguments: str | None = None + func_content_end = tool_text.find( + self.function_end_token, func_start + ) + if func_content_end != -1: + parsed_tool = self._parse_xml_function_call( + tool_text[func_start:func_content_end] + ) + if parsed_tool is not None: + complete_arguments = parsed_tool.function.arguments + # Always append — each tool call is a separate # invocation even if the function name is the same # (e.g. two consecutive "read" calls). self.prev_tool_call_arr.append( { "name": self.current_function_name, - "arguments": "{}", + "arguments": complete_arguments or "{}", } ) @@ -440,7 +743,12 @@ def extract_tool_calls_streaming( # compute remaining arguments at stream end. Without # this, IndexError occurs when the serving layer # accesses streamed_args_for_tool[index]. - self.streamed_args_for_tool.append("") + self.streamed_args_for_tool.append(complete_arguments or "") + + if complete_arguments is not None: + self.json_started = True + self.json_closed = True + self.in_function = False # Send header with function info return DeltaMessage( @@ -449,7 +757,8 @@ def extract_tool_calls_streaming( index=self.current_tool_index, id=self.current_tool_id, function=DeltaFunctionCall( - name=self.current_function_name, arguments="" + name=self.current_function_name, + arguments=complete_arguments or "", ), type="function", ) @@ -476,15 +785,10 @@ def extract_tool_calls_streaming( ] ) - # Find all parameter start positions in current tool_text - param_starts = [] - search_idx = 0 - while True: - search_idx = tool_text.find(self.parameter_prefix, search_idx) - if search_idx == -1: - break - param_starts.append(search_idx) - search_idx += len(self.parameter_prefix) + param_config = find_tool_properties( + self.tools, self.current_function_name or "" + ) + param_starts = self._parameter_start_positions(tool_text, param_config) # Process ALL complete params in a loop (spec decode fix). # With speculative decoding a single delta can deliver @@ -506,49 +810,47 @@ def extract_tool_calls_streaming( current_param_name = remaining[:name_end] value_start = param_start + name_end + 1 - value_text = tool_text[value_start:] - if value_text.startswith("\n"): - value_text = value_text[1:] - - param_end_idx = value_text.find(self.parameter_end_token) - if param_end_idx == -1: - next_param_idx = value_text.find(self.parameter_prefix) - func_end_idx = value_text.find(self.function_end_token) - - if next_param_idx != -1 and ( - func_end_idx == -1 or next_param_idx < func_end_idx - ): - param_end_idx = next_param_idx - elif func_end_idx != -1: - param_end_idx = func_end_idx + next_param_idx = ( + param_starts[self.param_count + 1] + if self.param_count + 1 < len(param_starts) + else None + ) + function_end_idx = tool_text.find(self.function_end_token, value_start) + tool_end_idx = tool_text.find(self.tool_call_end_token, value_start) + boundaries = [ + boundary + for boundary in (function_end_idx, tool_end_idx) + if boundary != -1 + ] + container_end = min(boundaries) if boundaries else None + value_end = self._find_structural_parameter_end( + tool_text, + value_start=value_start, + next_parameter_start=next_param_idx, + container_end=container_end, + allow_text_end=False, + ) + if value_end is None: + if next_param_idx is not None: + # Recover a missing before a real next tag. + value_end = next_param_idx + elif container_end is not None: + # Recover a missing final close before function/tool end. + value_end = container_end else: - # Fallback for malformed XML where - # is missing. Use as a delimiter - # if present in the value so we don't include - # the closing tag as part of the param value. - tool_end_in_value = value_text.find(self.tool_call_end_token) - if tool_end_in_value != -1: - param_end_idx = tool_end_in_value - else: - # Parameter incomplete — break so we still - # emit any fragments accumulated by earlier - # loop iterations. - break - - if param_end_idx == -1: - break + # Wait for structural lookahead. The current suffix may + # be literal text inside a long value. + break - param_value = value_text[:param_end_idx] + param_value = tool_text[value_start:value_end] + if param_value.startswith("\n"): + param_value = param_value[1:] if param_value.endswith("\n"): param_value = param_value[:-1] self.current_param_name = current_param_name self.accumulated_params[current_param_name] = param_value - param_config = find_tool_properties( - self.tools, self.current_function_name or "" - ) - converted_value = self._convert_param_value( param_value, current_param_name, @@ -578,6 +880,14 @@ def extract_tool_calls_streaming( len(self.streamed_args_for_tool), ) + # A speculative burst (or stream_interval > 1) can complete + # the final parameter and close the function in one parser + # call. Returning only the parameter fragment loses the final + # JSON brace because the next delta may be EOS/empty. + if not self.json_closed and self.function_end_token in tool_text: + self._finish_streaming_function(tool_text) + combined += "}" + return DeltaMessage( tool_calls=[ DeltaToolCall( @@ -594,39 +904,7 @@ def extract_tool_calls_streaming( # "}" and set in_function=False before the parameter loop # ever ran, causing the parameter to be silently dropped. if not self.json_closed and self.function_end_token in tool_text: - self.json_closed = True - - func_start = tool_text.find(self.tool_call_prefix) + len( - self.tool_call_prefix - ) - func_content_end = tool_text.find(self.function_end_token, func_start) - if func_content_end != -1: - func_content = tool_text[func_start:func_content_end] - try: - parsed_tool = self._parse_xml_function_call( - func_content, - ) - if parsed_tool and self.current_tool_index < len( - self.prev_tool_call_arr - ): - self.prev_tool_call_arr[self.current_tool_index][ - "arguments" - ] = parsed_tool.function.arguments - except Exception: - logger.debug( - "Failed to parse tool call during streaming: %s", - tool_text, - exc_info=True, - ) - - if self.current_tool_index < len(self.streamed_args_for_tool): - self.streamed_args_for_tool[self.current_tool_index] += "}" - else: - logger.warning( - "streamed_args_for_tool out of sync: index=%d len=%d", - self.current_tool_index, - len(self.streamed_args_for_tool), - ) + self._finish_streaming_function(tool_text) result = DeltaMessage( tool_calls=[ @@ -637,10 +915,6 @@ def extract_tool_calls_streaming( ] ) - self.in_function = False - self.json_closed = True - self.accumulated_params = {} - return result return None diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index f940739a96..227341a0cc 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -103,6 +103,9 @@ def __getitem__(self, key): fireredlid="FireRedLIDConfig", funaudiochat="FunAudioChatConfig", granite4_vision="Granite4VisionConfig", + glm5_next="Glm5NextConfig", + glm5_next_text="Glm5NextTextConfig", + glm5_next_vision="Glm5NextVisionConfig", hyperclovax="HyperCLOVAXConfig", hyperclovax_vlm="HCXVisionConfig", hunyuan_vl="HunYuanVLConfig", @@ -131,6 +134,8 @@ def __getitem__(self, key): qianfan_ocr="QianfanOCRConfig", qwen3_asr="Qwen3ASRConfig", qwen3_next="Qwen3NextConfig", + qwen4_exp="Qwen4ExpConfig", + qwen4_exp_text="Qwen4ExpTextConfig", qwen3_5="Qwen3_5Config", qwen3_5_moe="Qwen3_5MoeConfig", laguna="LagunaConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 5998e61dfd..16200da48a 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -35,6 +35,9 @@ "FunAudioChatConfig": "vllm.transformers_utils.configs.funaudiochat", "FunAudioChatAudioEncoderConfig": "vllm.transformers_utils.configs.funaudiochat", "Granite4VisionConfig": "vllm.transformers_utils.configs.granite4_vision", + "Glm5NextConfig": "vllm.transformers_utils.configs.glm5_next", + "Glm5NextTextConfig": "vllm.transformers_utils.configs.glm5_next", + "Glm5NextVisionConfig": "vllm.transformers_utils.configs.glm5_next", "HunYuanVLConfig": "vllm.transformers_utils.configs.hunyuan_vl", "HunYuanVLTextConfig": "vllm.transformers_utils.configs.hunyuan_vl", "HunYuanVLVisionConfig": "vllm.transformers_utils.configs.hunyuan_vl", @@ -75,6 +78,9 @@ "QianfanOCRVisionConfig": "vllm.transformers_utils.configs.qianfan_ocr", "Qwen3ASRConfig": "vllm.transformers_utils.configs.qwen3_asr", "Qwen3NextConfig": "vllm.transformers_utils.configs.qwen3_next", + "Qwen4ExpConfig": "vllm.transformers_utils.configs.qwen4_exp", + "Qwen4ExpTextConfig": "vllm.transformers_utils.configs.qwen4_exp", + "Qwen4ExpVisionConfig": "vllm.transformers_utils.configs.qwen4_exp", "Qwen3_5Config": "vllm.transformers_utils.configs.qwen3_5", "Qwen3_5TextConfig": "vllm.transformers_utils.configs.qwen3_5", "Qwen3_5MoeConfig": "vllm.transformers_utils.configs.qwen3_5_moe", @@ -106,6 +112,9 @@ "FunAudioChatConfig", "FunAudioChatAudioEncoderConfig", "Granite4VisionConfig", + "Glm5NextConfig", + "Glm5NextTextConfig", + "Glm5NextVisionConfig", "HunYuanVLConfig", "HunYuanVLTextConfig", "HunYuanVLVisionConfig", @@ -143,6 +152,9 @@ "QianfanOCRVisionConfig", "Qwen3ASRConfig", "Qwen3NextConfig", + "Qwen4ExpConfig", + "Qwen4ExpTextConfig", + "Qwen4ExpVisionConfig", "Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5MoeConfig", diff --git a/vllm/transformers_utils/configs/glm5_next.py b/vllm/transformers_utils/configs/glm5_next.py new file mode 100644 index 0000000000..29cae60584 --- /dev/null +++ b/vllm/transformers_utils/configs/glm5_next.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from transformers.configuration_utils import PretrainedConfig + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +class Glm5NextTextConfig(PretrainedConfig): + model_type = "glm5_next_text" + base_config_key = "text_config" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + model_type="glm5_next_text", + vocab_size: int = 154880, + hidden_size: int = 4096, + head_dim: int | None = None, + intermediate_size: int = 12288, + num_hidden_layers: int = 45, + num_attention_heads: int = 64, + num_key_value_heads: int | None = None, + hidden_act: str = "silu", + rms_norm_eps: float = 1e-5, + pad_token_id: int | None = 151329, + bos_token_id: int | None = None, + eos_token_id: int | list[int] | None = None, + rope_parameters: dict | None = None, + max_position_embeddings: int = 1048576, + tie_word_embeddings: bool = False, + moe_intermediate_size: int = 2048, + moe_renormalize: bool = True, + scoring_func: str = "sigmoid", + n_routed_experts: int | None = 288, + num_experts_per_token: int = 7, + n_shared_experts: int = 1, + routed_scaling_factor: float = 2.5, + first_k_dense_replace: int = 0, + moe_layer_freq: int = 1, + use_grouped_topk: bool = True, + n_group: int = 1, + topk_group: int = 1, + mla: bool = True, + q_lora_rank: int | None = 1536, + kv_lora_rank: int | None = 512, + qk_nope_head_dim: int = 256, + qk_rope_head_dim: int = 0, + v_head_dim: int | None = 256, + mla_nope: bool | None = True, + num_nextn_predict_layers: int = 1, + # Per-layer layout: "linear_attention" | "deepseek_sparse_attention" + layer_types: list[str] | None = None, + # Per-layer MLP: "dense" | "sparse" + mlp_layer_types: list[str] | None = None, + # Linear-attention (KDA) head config (flattened from the old + # linear_attn_config dict). + linear_head_dim: int = 128, + linear_num_heads: int = 64, + linear_conv_kernel_dim: int = 4, + linear_lower_bound: float | None = -5.0, + index_head_dim: int | None = None, + index_topk: int | None = None, + index_n_heads: int | None = None, + index_dsa_use_layernorm: bool = True, + index_kpool_compress: bool = True, + # Every ``index_kpool`` indexer K entries pool into one stored entry + # (compress_ratio). The 300B checkpoint ships 4; topk runs at pool + # granularity (select_k = index_topk // index_kpool). + index_kpool: int | None = 4, + index_kpool_always_select_tail: bool = True, + mhc: bool | None = True, + mhc_num_residual_streams: int = 4, + hc_eps: float | None = 1e-06, + mhc_tau: float = 0.05, + hres_vwnstyle: bool | None = True, + mhc_no_norm_weight: bool | None = False, + mhc_sinkhorn_iterations: int | None = 20, + mhc_post_mult_value: float | None = 2.0, + swiglu_limit: float | None = None, + **kwargs, + ): + # Preserve checkpoint field names and local aliases because their + # consumers use different spellings. + num_experts_per_token = kwargs.get("num_experts_per_tok", num_experts_per_token) + mhc_num_residual_streams = kwargs.get("hc_mult", mhc_num_residual_streams) + mhc_sinkhorn_iterations = kwargs.get( + "hc_sinkhorn_iters", mhc_sinkhorn_iterations + ) + # Checkpoint ships ``mla_use_nope`` (not ``mla_nope``); without this + # alias self.mla_nope silently stays at the param default. + mla_nope = kwargs.get("mla_use_nope", mla_nope) + # Checkpoints ship the KDA head config as the ``linear_attn_config`` + # dict (head_dim / num_heads / short_conv_kernel_size / + # gate_lower_bound) rather than the flattened top-level fields; fold it + # in so the trained values are read instead of the param defaults + # (which only match this checkpoint by coincidence). + linear_cfg = kwargs.get("linear_attn_config") or {} + if linear_cfg: + linear_head_dim = linear_cfg.get("head_dim", linear_head_dim) + linear_num_heads = linear_cfg.get("num_heads", linear_num_heads) + linear_conv_kernel_dim = linear_cfg.get( + "short_conv_kernel_size", linear_conv_kernel_dim + ) + linear_lower_bound = linear_cfg.get("gate_lower_bound", linear_lower_bound) + + self.model_type = model_type + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.head_dim = ( + head_dim if head_dim is not None else hidden_size // num_attention_heads + ) + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.rms_norm_eps = rms_norm_eps + self.max_position_embeddings = max_position_embeddings + self.rope_parameters = rope_parameters + + # mla config + self.mla = mla + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.mla_nope = mla_nope + # moe config + self.n_routed_experts = n_routed_experts + self.num_experts_per_token = num_experts_per_token + self.moe_renormalize = moe_renormalize + self.n_shared_experts = n_shared_experts + self.routed_scaling_factor = routed_scaling_factor + self.scoring_func = scoring_func + assert self.scoring_func in ("softmax", "sigmoid") + self.moe_intermediate_size = moe_intermediate_size + self.first_k_dense_replace = first_k_dense_replace + self.moe_layer_freq = moe_layer_freq + self.use_grouped_topk = use_grouped_topk + self.n_group = n_group + self.topk_group = topk_group + self.num_nextn_predict_layers = num_nextn_predict_layers + + # Per-layer attention / MLP layout. Normalize mlp_layer_types from + # first_k_dense_replace when the new-schema field is absent so layer + # construction sees a consistent layout (mirrors cohere2_moe). + self.layer_types = layer_types + if mlp_layer_types is None: + n = self.num_hidden_layers + if first_k_dense_replace is not None: + mlp_layer_types = ["dense"] * first_k_dense_replace + ["sparse"] * ( + n - first_k_dense_replace + ) + else: + mlp_layer_types = ["sparse"] * n + self.mlp_layer_types = mlp_layer_types + + # Linear-attention (KDA) head config. + self.linear_head_dim = linear_head_dim + self.linear_num_heads = linear_num_heads + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_lower_bound = linear_lower_bound + + # dsa index config + self.index_head_dim = index_head_dim + self.index_topk = index_topk + self.index_n_heads = index_n_heads + self.index_dsa_use_layernorm = index_dsa_use_layernorm + self.index_kpool_compress = index_kpool_compress + self.index_kpool = index_kpool + self.index_kpool_always_select_tail = index_kpool_always_select_tail + + # mhc config + self.mhc = mhc + self.mhc_num_residual_streams = mhc_num_residual_streams + self.mhc_tau = mhc_tau + self.hres_vwnstyle = hres_vwnstyle + self.hc_eps = hc_eps + self.mhc_no_norm_weight = mhc_no_norm_weight + self.mhc_sinkhorn_iterations = mhc_sinkhorn_iterations + self.mhc_post_mult_value = mhc_post_mult_value + + self.swiglu_limit = swiglu_limit + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @property + def is_mla(self): + return ( + self.q_lora_rank is not None + or self.kv_lora_rank is not None + or self.qk_nope_head_dim is not None + or self.qk_rope_head_dim is not None + or self.v_head_dim is not None + or self.mla_nope is True + ) + + @property + def is_moe(self): + return self.n_routed_experts is not None + + @property + def is_linear_attn(self) -> bool: + return self.layer_types is not None and any( + t == "linear_attention" for t in self.layer_types + ) + + def is_kda_layer(self, layer_idx: int): + return ( + self.layer_types is not None + and layer_idx < len(self.layer_types) + and self.layer_types[layer_idx] == "linear_attention" + ) + + @property + def layers_block_type(self): + # Map the schema's per-layer types onto the block strings vLLM's hybrid + # accounting (get_num_layers_by_block_type) recognizes: linear-attention + # layers stay "linear_attention"; every other attention variant collapses + # to "attention". + if self.layer_types is None: + return ["attention"] * self.num_hidden_layers + return [ + "linear_attention" if t == "linear_attention" else "attention" + for t in self.layer_types + ] + + +class Glm5NextVisionConfig(PretrainedConfig): + model_type = "glm5_next_vision" + base_config_key = "vision_config" + + def __init__( + self, + depth: int = 24, + hidden_size: int = 1024, + hidden_act: str = "silu", + image_size: int = 448, + intermediate_size: int = 4096, + num_heads: int = 16, + out_hidden_size: int = 4096, + projection_intermediate_size: int = 10240, + in_channels: int = 3, + initializer_range: float = 0.02, + patch_size: int = 14, + rms_norm_eps: float = 1e-5, + spatial_merge_size: int = 2, + temporal_patch_size: int = 2, + attention_dropout: float = 0.0, + attention_bias: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + + self.depth = depth + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.image_size = image_size + self.intermediate_size = intermediate_size + self.num_heads = num_heads + self.out_hidden_size = out_hidden_size + # GLM-5.3-Flash merger bottleneck width (absent from the generic + # GLM-OCR vision config); the tower uses it as the PatchMerger + # context_dim instead of text_config.intermediate_size. + self.projection_intermediate_size = projection_intermediate_size + self.in_channels = in_channels + self.initializer_range = initializer_range + self.patch_size = patch_size + # GLM-5.3-Flash checkpoints ship vision_config.rms_norm_eps = 1e-5, + # but the vision tower was trained with 1e-6. Serving with 1e-5 drifts + # the RMSNorm and produces repetitive/degraded image descriptions, so + # force the trained value regardless of the checkpoint field. + self.rms_norm_eps = 1e-6 + self.spatial_merge_size = spatial_merge_size + self.temporal_patch_size = temporal_patch_size + self.attention_dropout = attention_dropout + self.attention_bias = attention_bias + + +class Glm5NextConfig(PretrainedConfig): + model_type = "glm5_next" + sub_configs = { + "vision_config": Glm5NextVisionConfig, + "text_config": Glm5NextTextConfig, + } + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config=None, + vision_config=None, + image_token_id: int = 154854, + video_token_id: int = 154855, + image_start_token_id: int = 154830, + image_end_token_id: int = 154831, + video_start_token_id: int = 154832, + video_end_token_id: int = 154833, + **kwargs, + ): + # Init super() first so base-class defaults don't clobber text-config + # values set below (PretrainedConfig has many text-related defaults + # that differ from Glm5NextTextConfig). + super().__init__(**kwargs) + + if isinstance(vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**vision_config) + elif vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + + if isinstance(text_config, dict): + self.text_config = self.sub_configs["text_config"](**text_config) + elif text_config is None: + # Backward compatibility: a flat top-level checkpoint (no nested + # text_config) folds its text fields into Glm5NextTextConfig. + self.text_config = self.sub_configs["text_config"](**kwargs) + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.image_start_token_id = image_start_token_id + self.image_end_token_id = image_end_token_id + self.video_start_token_id = video_start_token_id + self.video_end_token_id = video_end_token_id + + # Mirror attention implementation recursively onto sub-configs. + self._attn_implementation = kwargs.pop("attn_implementation", None) + + # Config-metadata fields that belong to the top-level (multimodal) config + # and must NOT be mirrored onto text_config: ``architectures`` / + # ``torch_dtype`` differ between the top-level config and the text + # sub-config, and mirroring them makes the top-level ``architectures`` + # silently read back as None (PretrainedConfig initializes both to None), + # which then fails model-class resolution ("No model architectures are + # specified"). + _UNMIRRORED_KEYS = [ + "_name_or_path", + "model_type", + "dtype", + "torch_dtype", + "architectures", + "_attn_implementation_internal", + ] + + def __setattr__(self, key, value): + unmirrored = type(self)._UNMIRRORED_KEYS + if ( + (text_config := super().__getattribute__("__dict__").get("text_config")) + is not None + and key not in unmirrored + and key in text_config.__dict__ + ): + setattr(text_config, key, value) + else: + super().__setattr__(key, value) + + def __getattribute__(self, key): + unmirrored = type(self)._UNMIRRORED_KEYS + if ( + "text_config" in super().__getattribute__("__dict__") + and key not in unmirrored + ): + text_config = super().__getattribute__("text_config") + # Forward both instance attributes AND class-defined properties/ + # methods of the text config, so a flat text-only checkpoint + # (model_type "glm5_next", no nested text_config) sees is_moe / + # is_kda_layer / layers_block_type like a Glm5NextTextConfig. + if key in text_config.__dict__ or key in type(text_config).__dict__: + return getattr(text_config, key) + + return super().__getattribute__(key) diff --git a/vllm/transformers_utils/configs/qwen4_exp.py b/vllm/transformers_utils/configs/qwen4_exp.py new file mode 100644 index 0000000000..ed4c6cca8e --- /dev/null +++ b/vllm/transformers_utils/configs/qwen4_exp.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exports for the package-owned Qwen4Exp configs.""" + +from vllm.models.qwen4_exp.config import ( + Qwen4ExpConfig, + Qwen4ExpTextConfig, + Qwen4ExpVisionConfig, +) + +__all__ = [ + "Qwen4ExpConfig", + "Qwen4ExpTextConfig", + "Qwen4ExpVisionConfig", +] diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 8545219753..d33169d03d 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -227,6 +227,8 @@ def is_deepseek_mla(self) -> bool: "deepseek_v32", "deepseek_v4", "deepseek_mtp", + "glm5_next", + "glm5_next_text", "glm_moe_dsa", "glm4_moe_lite", "glm4_moe_lite_mtp", diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index ba2872f892..09cd5a53bc 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -19,6 +19,7 @@ "FireRedLIDProcessor", "FunASRProcessor", "GLM4VProcessor", + "Glm5NextProcessor", "Granite4VisionProcessor", "H2OVLProcessor", "HunYuanVLProcessor", @@ -52,6 +53,7 @@ "FireRedLIDProcessor": "vllm.transformers_utils.processors.fireredlid", "FunASRProcessor": "vllm.transformers_utils.processors.funasr", "GLM4VProcessor": "vllm.transformers_utils.processors.glm4v", + "Glm5NextProcessor": "vllm.transformers_utils.processors.glm5next", "Granite4VisionProcessor": "vllm.transformers_utils.processors.granite4_vision", "H2OVLProcessor": "vllm.transformers_utils.processors.h2ovl", "HunYuanVLProcessor": "vllm.transformers_utils.processors.hunyuan_vl", diff --git a/vllm/transformers_utils/processors/glm5next.py b/vllm/transformers_utils/processors/glm5next.py new file mode 100644 index 0000000000..67b5446bc6 --- /dev/null +++ b/vllm/transformers_utils/processors/glm5next.py @@ -0,0 +1,1008 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM-native multimodal processor for GLM-5.3-Flash.""" + +import json +import math + +import numpy as np +import torch +from torchvision.transforms.v2 import functional as tvF +from transformers.image_processing_utils import BatchFeature +from transformers.image_processing_utils_fast import ( + BaseImageProcessorFast, + group_images_by_shape, + reorder_images, +) +from transformers.image_utils import ( + OPENAI_CLIP_MEAN, + OPENAI_CLIP_STD, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size, +) +from transformers.models.auto.image_processing_auto import get_image_processor_config +from transformers.processing_utils import ( + ImagesKwargs, + MultiModalData, + ProcessingKwargs, + ProcessorMixin, + Unpack, + VideosKwargs, +) +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput +from transformers.utils import TensorType, cached_file, logging +from transformers.video_processing_utils import BaseVideoProcessor +from transformers.video_utils import ( + VideoInput, + VideoMetadata, + group_videos_by_shape, + reorder_videos, +) + +logger = logging.get_logger(__name__) + +# Cap video inputs at 30,000 vision tokens to keep encoder profiling from +# starving the KV cache. Image inputs retain their checkpoint-defined budget. +_MAX_VIDEO_TOKENS = 30000 + +# Frame-sampler fallbacks (mirror the checkpoint's fps_interval=2 / +# max_frame_count_dynamic=2048); used when neither the request nor the +# processor config overrides them. +GLM_VIDEO_DEFAULT_FPS = 2.0 +GLM_VIDEO_DEFAULT_MAX_FRAMES = 2048 + + +def glm_sample_frame_indices( + total_frames: int, + fps: float, + duration: float, + *, + target_fps: float | None = None, + max_frame_count: int | None = None, + temporal_patch_size: int = 2, +) -> list[int]: + """GLM video frame sampling (training-reference parity). + + ``target_fps`` is the ``fps_interval`` request knob. The greedy walk + advances at ``1 / (temporal_patch_size * target_fps)`` seconds, so on + frame-dense sources it collects more candidates than ``extract_t`` and + the ``> extract_t`` fixup re-spreads the picks uniformly with + ``np.linspace`` -- that fallback is the intended reference behavior, not + an accident. Short clips (fewer frames than ``extract_t``) are spread at + evenly spaced timestamps (``floor`` sampling; the linspace variant + samples frames unevenly and cost 4 points on video grounding evals). + Request overrides: ``target_fps`` -> fps interval, ``max_frame_count`` + -> frame cap. + """ + max_frame_idx = total_frames - 1 + if not duration: + duration = (round(max_frame_idx / fps) + 1) if fps else 0 + if max_frame_count is None: + max_frame_count = GLM_VIDEO_DEFAULT_MAX_FRAMES + if target_fps is None: + target_fps = GLM_VIDEO_DEFAULT_FPS + + extract_t = int(duration * target_fps) + extract_t = min(extract_t, int(max_frame_count)) + + duration_per_frame = 1 / fps + timestamps = [i * duration_per_frame for i in range(total_frames)] + max_second = int(duration) + + if total_frames < extract_t: + frame_indices = [ + math.floor(_i * total_frames / extract_t) for _i in range(extract_t) + ] + else: + frame_indices = [] + current_second = 0.0 + inv_fps = 1 / (temporal_patch_size * target_fps) + for frame_index in range(total_frames): + if timestamps[frame_index] >= current_second: + current_second += inv_fps + frame_indices.append(frame_index) + if current_second >= max_second: + break + + if len(frame_indices) < extract_t: + if len(frame_indices) == 0: + start, end = 0, max(total_frames - 1, 0) + else: + start, end = frame_indices[0], frame_indices[-1] + frame_indices = np.linspace(start, end, extract_t, dtype=int).tolist() + elif len(frame_indices) > extract_t: + frame_indices = np.linspace(0, total_frames - 1, extract_t, dtype=int).tolist() + + seen, uniq = set(), [] + for idx in frame_indices: + if idx not in seen: + seen.add(idx) + uniq.append(int(idx)) + + if len(uniq) & 1: + uniq.append(uniq[-1]) + + return uniq + + +def _ceil_to_factor(value: int, factor: int) -> int: + """Round a positive integer upward to the nearest multiple of factor.""" + return math.ceil(value / factor) * factor + + +def _fit_aligned_size_within_budget( + t: int, + h: int, + w: int, + h_factor: int, + w_factor: int, + max_pixels: int, +) -> tuple[int, int]: + """Largest proportional size whose upward-aligned canvas fits the budget. + + Binary search on the unaligned content height; each candidate is rounded + upward to h_factor/w_factor, so the returned canvas always satisfies + ``t * aligned_h * aligned_w <= max_pixels``. + """ + minimum_pixels = t * h_factor * w_factor + if max_pixels < minimum_pixels: + raise ValueError( + f"max_pixels={max_pixels} is too small. At least " + f"{minimum_pixels} pixels are required for one aligned patch." + ) + + low, high = 1, h + best_h, best_w = h_factor, w_factor + while low <= high: + content_h = (low + high) // 2 + content_w = max(1, math.floor(w * content_h / h)) + aligned_h = _ceil_to_factor(content_h, h_factor) + aligned_w = _ceil_to_factor(content_w, w_factor) + if t * aligned_h * aligned_w <= max_pixels: + best_h, best_w = aligned_h, aligned_w + low = content_h + 1 + else: + high = content_h - 1 + return best_h, best_w + + +def smart_resize( + t: int, + h: int, + w: int, + t_factor: int = 1, + h_factor: int = 28, + w_factor: int = 28, + min_pixels: int = 56 * 56, + max_pixels: int = 14 * 14 * 4 * 1280, +) -> tuple[int, int]: + """GLM-5.3-Flash ``smart_resize``: upward-aligned canvas under a + ``t_bar * h_bar * w_bar`` pixel budget. + + Height/width always round UP to their factors (content is then padded, + never cropped or distorted); an over-budget canvas is refit by binary + search instead of one-shot square-root scaling. ``h_factor`` / + ``w_factor`` carry ``patch_expand_factor`` on top of + ``patch_size * merge_size``; ``t_factor`` is ``temporal_patch_size``. For + a still image ``t = t_factor = temporal_patch_size`` so ``t_bar = + temporal_patch_size``. + """ + if min(t, h, w, t_factor, h_factor, w_factor) <= 0: + raise ValueError("Image dimensions and alignment factors must be positive.") + if min_pixels <= 0 or max_pixels <= 0: + raise ValueError("min_pixels and max_pixels must be positive.") + if min_pixels > max_pixels: + raise ValueError("min_pixels must be less than or equal to max_pixels.") + + t_bar = max(t_factor, round(t / t_factor) * t_factor) + h_bar = _ceil_to_factor(h, h_factor) + w_bar = _ceil_to_factor(w, w_factor) + + if t_bar * h_bar * w_bar > max_pixels: + h_bar, w_bar = _fit_aligned_size_within_budget( + t=t_bar, + h=h, + w=w, + h_factor=h_factor, + w_factor=w_factor, + max_pixels=max_pixels, + ) + elif t_bar * h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (t * h * w)) + h_bar = _ceil_to_factor(max(1, math.ceil(h * beta)), h_factor) + w_bar = _ceil_to_factor(max(1, math.ceil(w * beta)), w_factor) + + # Alignment can push a candidate slightly over a tight max_pixels + # budget. Refit it when that happens. + if t_bar * h_bar * w_bar > max_pixels: + h_bar, w_bar = _fit_aligned_size_within_budget( + t=t_bar, + h=h, + w=w, + h_factor=h_factor, + w_factor=w_factor, + max_pixels=max_pixels, + ) + + return h_bar, w_bar + + +def _get_pad_content_size( + image_height: int, + image_width: int, + canvas_height: int, + canvas_width: int, + allow_upscale: bool = False, +) -> tuple[int, int]: + """Aspect-ratio-preserving content size that fits the canvas. + + Oversized images are shrunk proportionally. Small images are enlarged + only when ``allow_upscale``. Padding is applied after the resize. + """ + scale = min(canvas_height / image_height, canvas_width / image_width) + if not allow_upscale: + scale = min(1.0, scale) + content_height = max(1, min(canvas_height, math.floor(image_height * scale))) + content_width = max(1, min(canvas_width, math.floor(image_width * scale))) + return content_height, content_width + + +def _resize_or_pad( + stacked_images: torch.Tensor, + target_height: int, + target_width: int, + resize_mode: str, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + resize, + allow_upscale: bool = False, +) -> torch.Tensor: + """Resize onto the aligned canvas, or keep the aspect ratio and + zero-pad the right/bottom sides (``resize_mode="pad"``).""" + height, width = stacked_images.shape[-2:] + + if resize_mode == "resize": + return resize( + stacked_images, + size=SizeDict(height=target_height, width=target_width), + resample=resample, + ) + + if resize_mode != "pad": + raise ValueError("resize_mode must be either 'resize' or 'pad'.") + + content_height, content_width = _get_pad_content_size( + image_height=height, + image_width=width, + canvas_height=target_height, + canvas_width=target_width, + allow_upscale=allow_upscale, + ) + + if (content_height, content_width) != (height, width): + stacked_images = resize( + stacked_images, + size=SizeDict(height=content_height, width=content_width), + resample=resample, + ) + + # torchvision padding order: [left, top, right, bottom] -> pad only the + # right and bottom sides. + return tvF.pad( + stacked_images, + padding=[0, 0, target_width - content_width, target_height - content_height], + fill=0, + ) + + +def _pixel_budget( + min_image_tokens: int | None, + max_image_tokens: int | None, + patch_size: int, + merge_size: int, + temporal_patch_size: int, +) -> tuple[int, int]: + """(min_pixels, max_pixels) from the token bounds of + ``processor_config.json``; one vision token covers + ``temporal_patch_size * (patch_size * merge_size) ** 2`` pixels.""" + if min_image_tokens is None or max_image_tokens is None: + raise ValueError( + "min_image_tokens and max_image_tokens must be provided by " + "processor_config.json (or per-call kwargs)." + ) + factor = temporal_patch_size * (patch_size * merge_size) ** 2 + return min_image_tokens * factor, max_image_tokens * factor + + +class Glm5NextImageProcessorKwargs(ImagesKwargs, total=False): # type: ignore[call-arg] + patch_size: int | None + temporal_patch_size: int | None + merge_size: int | None + patch_expand_factor: int | None + resize_mode: str | None + min_image_tokens: int | None + max_image_tokens: int | None + + +class Glm5NextImageProcessor(BaseImageProcessorFast): + """Fast torchvision image processor for GLM-5.3-Flash. + + ``patch_expand_factor`` multiplies into the ``smart_resize`` spatial + factor ``patch_size * merge_size``. ``resize_mode`` picks the geometry: + ``"pad"`` (default) preserves the aspect ratio and zero-pads the + right/bottom of the upward-aligned canvas, ``"resize"`` stretches onto + it. Defaults mirror the checkpoint's ``image_processor`` config. + """ + + do_resize = True + resample = PILImageResampling.BICUBIC + size = {"longest_edge": 1} # unused: budgets come from the token bounds + do_rescale = True + do_normalize = True + image_mean = OPENAI_CLIP_MEAN + image_std = OPENAI_CLIP_STD + do_convert_rgb = True + patch_size = 14 + temporal_patch_size = 2 + merge_size = 2 + patch_expand_factor = 1 + resize_mode = "pad" + min_image_tokens = 16 + max_image_tokens = 8000 + valid_kwargs = Glm5NextImageProcessorKwargs + model_input_names = ["pixel_values", "image_grid_thw"] + + def _preprocess( + self, + images: list[torch.Tensor], + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + patch_size: int, + temporal_patch_size: int, + merge_size: int, + patch_expand_factor: int, + resize_mode: str | None, + min_image_tokens: int | None, + max_image_tokens: int | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + resize_mode = resize_mode if resize_mode is not None else self.resize_mode + min_pixels, max_pixels = _pixel_budget( + min_image_tokens if min_image_tokens is not None else self.min_image_tokens, + max_image_tokens if max_image_tokens is not None else self.max_image_tokens, + patch_size, + merge_size, + temporal_patch_size, + ) + grouped_images, grouped_images_index = group_images_by_shape( + images, disable_grouping=disable_grouping + ) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + height, width = stacked_images.shape[-2:] + if do_resize: + resized_height, resized_width = smart_resize( + t=temporal_patch_size, + h=height, + w=width, + t_factor=temporal_patch_size, + h_factor=patch_size * merge_size * patch_expand_factor, + w_factor=patch_size * merge_size * patch_expand_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + stacked_images = _resize_or_pad( + stacked_images, + target_height=resized_height, + target_width=resized_width, + resize_mode=resize_mode, + resample=resample, + resize=self.resize, + allow_upscale=(temporal_patch_size * height * width < min_pixels), + ) + resized_images_grouped[shape] = stacked_images + + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + grouped_images, grouped_images_index = group_images_by_shape( + resized_images, disable_grouping=disable_grouping + ) + processed_images_grouped = {} + processed_grids = {} + + for shape, stacked_images in grouped_images.items(): + resized_height, resized_width = stacked_images.shape[-2:] + + patches = self.rescale_and_normalize( + stacked_images, + do_rescale, + rescale_factor, + do_normalize, + image_mean, + image_std, + ) + if patches.ndim == 4: # (B, C, H, W) + patches = patches.unsqueeze(1) # (B, T=1, C, H, W) + + if patches.shape[1] % temporal_patch_size != 0: + repeats = patches[:, -1:].repeat( + 1, + temporal_patch_size - (patches.shape[1] % temporal_patch_size), + 1, + 1, + 1, + ) + patches = torch.cat([patches, repeats], dim=1) + + batch_size, t_len, channel = patches.shape[:3] + grid_t = t_len // temporal_patch_size + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + + patches = patches.view( + batch_size, + grid_t, + temporal_patch_size, + channel, + grid_h // merge_size, + merge_size, + patch_size, + grid_w // merge_size, + merge_size, + patch_size, + ) + # (B, grid_t, gh, gw, mh, mw, C, tp, ph, pw) + patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9) + + flatten_patches = patches.reshape( + batch_size, + grid_t * grid_h * grid_w, + channel * temporal_patch_size * patch_size * patch_size, + ) + + processed_images_grouped[shape] = flatten_patches + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size + + processed_images = reorder_images( + processed_images_grouped, grouped_images_index + ) + processed_grids = reorder_images(processed_grids, grouped_images_index) + + pixel_values = torch.cat(processed_images, dim=0) + image_grid_thw = torch.tensor(processed_grids) + + return BatchFeature( + data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, + tensor_type=return_tensors, + ) + + def preprocess( + self, images: ImageInput, **kwargs: Unpack[Glm5NextImageProcessorKwargs] + ) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def get_number_of_image_patches( + self, height: int, width: int, images_kwargs: dict | None = None + ) -> int: + """Number of image patches (pre-merge) for a given (height, width).""" + images_kwargs = images_kwargs or {} + patch_size = images_kwargs.get("patch_size", self.patch_size) + merge_size = images_kwargs.get("merge_size", self.merge_size) + patch_expand_factor = images_kwargs.get( + "patch_expand_factor", self.patch_expand_factor + ) + min_pixels, max_pixels = _pixel_budget( + images_kwargs.get("min_image_tokens", self.min_image_tokens), + images_kwargs.get("max_image_tokens", self.max_image_tokens), + patch_size, + merge_size, + self.temporal_patch_size, + ) + resized_height, resized_width = smart_resize( + t=self.temporal_patch_size, + h=height, + w=width, + t_factor=self.temporal_patch_size, + h_factor=patch_size * merge_size * patch_expand_factor, + w_factor=patch_size * merge_size * patch_expand_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + return grid_h * grid_w + + +class Glm5NextVideoProcessorKwargs(VideosKwargs, total=False): # type: ignore[call-arg] + fps: list[float] | float + patch_size: int + temporal_patch_size: int + merge_size: int + patch_expand_factor: int + resize_mode: str | None + target_fps: float | None + max_frames: int | None + fps_interval: int | None + max_frame_count_dynamic: int | None + min_image_tokens: int | None + max_image_tokens: int | None + + +class Glm5NextVideoProcessor(BaseVideoProcessor): + """Fast video processor for GLM-5.3-Flash. + + Shares ``smart_resize`` / the pad-mode geometry / the patchify with the + image processor, and adds GLM-5.3-Flash frame sampling + (``glm_sample_frame_indices``: ``fps_interval`` semantics with a + temporal-patch-scaled greedy walk). Defaults mirror the checkpoint's + ``video_processor`` config. + """ + + resample = PILImageResampling.BICUBIC + size = {"longest_edge": 1} # unused: budgets come from the token bounds + image_mean = OPENAI_CLIP_MEAN + image_std = OPENAI_CLIP_STD + do_resize = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + do_sample_frames = True + patch_size = 14 + temporal_patch_size = 2 + patch_expand_factor = 1 + merge_size = 2 + valid_kwargs = Glm5NextVideoProcessorKwargs + num_frames = 16 + fps = 2 + fps_interval = 2.0 + max_frame_count_dynamic = 2048 + resize_mode = "pad" + min_image_tokens = 16 + max_image_tokens = 240000 + model_input_names = ["pixel_values_videos", "video_grid_thw"] + + def sample_frames( + self, + metadata: VideoMetadata, + fps: int | float | None = None, + **kwargs, + ) -> np.ndarray: + """Sample frame indices with GLM's fps-interval policy. + + ``fps`` / ``target_fps``, ``max_frames`` and ``fps_interval`` / + ``max_frame_count_dynamic`` are the overrides described in + :func:`glm_sample_frame_indices`. + """ + if metadata is None or getattr(metadata, "fps", None) is None: + raise ValueError( + "Asked to sample frames per second but no video metadata was " + "provided which is required when sampling in GLM-5.3-Flash. Please " + "pass in `VideoMetadata` object or set `do_sample_frames=False`." + ) + + target_fps = fps if fps is not None else kwargs.get("target_fps") + if target_fps is None: + target_fps = self.fps_interval + indices = glm_sample_frame_indices( + metadata.total_num_frames, + metadata.fps, + metadata.duration or 0, + target_fps=target_fps, + max_frame_count=kwargs.get("max_frames") or self.max_frame_count_dynamic, + temporal_patch_size=self.temporal_patch_size, + ) + return np.array(indices) + + def _preprocess( + self, + videos: list[torch.Tensor], + do_convert_rgb: bool = True, + do_resize: bool = True, + size: SizeDict | None = None, + resample: "PILImageResampling | int | None" = PILImageResampling.BICUBIC, + do_rescale: bool = True, + rescale_factor: float = 1 / 255.0, + do_normalize: bool = True, + image_mean: float | list[float] | None = None, + image_std: float | list[float] | None = None, + patch_size: int | None = None, + temporal_patch_size: int | None = None, + patch_expand_factor: int | None = None, + merge_size: int | None = None, + resize_mode: str | None = None, + min_image_tokens: int | None = None, + max_image_tokens: int | None = None, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchFeature: + patch_expand_factor = self.patch_expand_factor + patch_size = patch_size if patch_size is not None else self.patch_size + temporal_patch_size = ( + temporal_patch_size + if temporal_patch_size is not None + else self.temporal_patch_size + ) + merge_size = merge_size if merge_size is not None else self.merge_size + resize_mode = resize_mode if resize_mode is not None else self.resize_mode + min_pixels, max_pixels = _pixel_budget( + min_image_tokens if min_image_tokens is not None else self.min_image_tokens, + max_image_tokens if max_image_tokens is not None else self.max_image_tokens, + patch_size, + merge_size, + temporal_patch_size, + ) + grouped_videos, grouped_videos_index = group_videos_by_shape(videos) + resized_videos_grouped = {} + for shape, stacked_videos in grouped_videos.items(): + if do_convert_rgb: + stacked_videos = self.convert_to_rgb(stacked_videos) + b, t_len, c, h, w = stacked_videos.shape + num_frames, height, width = t_len, h, w + if do_resize: + resized_height, resized_width = smart_resize( + t=num_frames, + h=height, + w=width, + t_factor=temporal_patch_size, + h_factor=patch_size * merge_size * patch_expand_factor, + w_factor=patch_size * merge_size * patch_expand_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + stacked_videos = stacked_videos.view(b * t_len, c, h, w) + stacked_videos = _resize_or_pad( + stacked_videos, + target_height=resized_height, + target_width=resized_width, + resize_mode=resize_mode, + resample=resample, + resize=self.resize, + allow_upscale=(num_frames * height * width < min_pixels), + ) + stacked_videos = stacked_videos.view( + b, t_len, c, resized_height, resized_width + ) + resized_videos_grouped[shape] = stacked_videos + resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index) + + grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos) + processed_videos_grouped = {} + processed_grids = {} + for shape, stacked_videos in grouped_videos.items(): + resized_height, resized_width = get_image_size( + stacked_videos[0], channel_dim=ChannelDimension.FIRST + ) + stacked_videos = self.rescale_and_normalize( + stacked_videos, + do_rescale, + rescale_factor, + do_normalize, + image_mean, + image_std, + ) + patches = stacked_videos + + if pad := -patches.shape[1] % temporal_patch_size: + repeats = patches[:, -1:].expand(-1, pad, -1, -1, -1) + patches = torch.cat((patches, repeats), dim=1) + batch_size, grid_t, channel = patches.shape[:3] + grid_t = grid_t // temporal_patch_size + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + + patches = patches.view( + batch_size, + grid_t, + temporal_patch_size, + channel, + grid_h // merge_size, + merge_size, + patch_size, + grid_w // merge_size, + merge_size, + patch_size, + ) + patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9) + flatten_patches = patches.reshape( + batch_size, + grid_t * grid_h * grid_w, + channel * temporal_patch_size * patch_size * patch_size, + ) + + processed_videos_grouped[shape] = flatten_patches + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size + + processed_videos = reorder_videos( + processed_videos_grouped, grouped_videos_index + ) + processed_grids = reorder_videos(processed_grids, grouped_videos_index) + pixel_values_videos = torch.cat(processed_videos, dim=0) + video_grid_thw = torch.tensor(processed_grids) + return BatchFeature( + data={ + "pixel_values_videos": pixel_values_videos, + "video_grid_thw": video_grid_thw, + }, + tensor_type=return_tensors, + ) + + def get_number_of_video_patches( + self, + num_frames: int, + height: int, + width: int, + videos_kwargs: dict | None = None, + ) -> int: + """Number of temporal/spatial patches for a raw video shape.""" + videos_kwargs = videos_kwargs or {} + patch_size = videos_kwargs.get("patch_size", self.patch_size) + temporal_patch_size = videos_kwargs.get( + "temporal_patch_size", self.temporal_patch_size + ) + merge_size = videos_kwargs.get("merge_size", self.merge_size) + patch_expand_factor = videos_kwargs.get( + "patch_expand_factor", self.patch_expand_factor + ) + min_pixels, max_pixels = _pixel_budget( + videos_kwargs.get("min_image_tokens", self.min_image_tokens), + videos_kwargs.get("max_image_tokens", self.max_image_tokens), + patch_size, + merge_size, + temporal_patch_size, + ) + resized_height, resized_width = smart_resize( + t=max(num_frames, temporal_patch_size), + h=height, + w=width, + t_factor=temporal_patch_size, + h_factor=patch_size * merge_size * patch_expand_factor, + w_factor=patch_size * merge_size * patch_expand_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + padded_frames = num_frames + (-num_frames % temporal_patch_size) + grid_t = max(padded_frames // temporal_patch_size, 1) + grid_h = resized_height // patch_size + grid_w = resized_width // patch_size + return grid_t * grid_h * grid_w + + +class Glm5NextProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-arg] + images_kwargs: Glm5NextImageProcessorKwargs + videos_kwargs: Glm5NextVideoProcessorKwargs + _defaults = { + "text_kwargs": { + "padding": False, + "return_token_type_ids": False, + "return_mm_token_type_ids": False, + }, + "videos_kwargs": {"return_metadata": True}, + } + + +class Glm5NextProcessor(ProcessorMixin): + """Wraps a GLM-5.3-Flash image processor, video processor and tokenizer. + + Token expansion per image = ``prod(image_grid_thw) // merge_size**2``; video + frames are expanded with ``<|begin_of_image|>...<|end_of_image|>{ts} seconds`` + structure (mrope timestamps). + """ + + attributes = ["image_processor", "tokenizer", "video_processor"] + image_processor_class = "AutoImageProcessor" + video_processor_class = "AutoVideoProcessor" + tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast") + + def __init__( + self, + image_processor=None, + tokenizer=None, + video_processor=None, + chat_template=None, + **kwargs, + ) -> None: + super().__init__( + image_processor, tokenizer, video_processor, chat_template=chat_template + ) + self.image_token = ( + "<|image|>" + if not hasattr(tokenizer, "image_token") + else tokenizer.image_token + ) + self.video_token = ( + "<|video|>" + if not hasattr(tokenizer, "video_token") + else tokenizer.video_token + ) + self.image_token_id = ( + tokenizer.image_token_id + if getattr(tokenizer, "image_token_id", None) + else tokenizer.convert_tokens_to_ids(self.image_token) + ) + self.video_token_id = ( + tokenizer.video_token_id + if getattr(tokenizer, "video_token_id", None) + else tokenizer.convert_tokens_to_ids(self.video_token) + ) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + """Build the processor directly from the checkpoint config. + + GLM-5.3-Flash stores nested image/video configs in + ``processor_config.json`` and declares a custom processor class. This + method reads those configs directly and caps only the video token budget. + """ + from transformers import AutoTokenizer + + model_path = pretrained_model_name_or_path + tokenizer = AutoTokenizer.from_pretrained(model_path, **kwargs) + + def _cap_cfg(cfg: dict, *, is_video: bool) -> dict: + # Video keeps a serving token cap (the checkpoint's 240k-token + # budget would starve the KV cache at startup profiling); images + # follow the checkpoint budget verbatim so preprocessing matches + # the HF reference exactly. + if is_video and cfg.get("max_image_tokens") is not None: + cfg["max_image_tokens"] = min( + cfg["max_image_tokens"], _MAX_VIDEO_TOKENS + ) + return cfg + + ip_cfg = _cap_cfg(dict(get_image_processor_config(model_path)), is_video=False) + image_processor = Glm5NextImageProcessor( + **{k: v for k, v in ip_cfg.items() if k != "image_processor_type"} + ) + + cached_file_kwargs = { + key: kwargs[key] + for key in ( + "cache_dir", + "force_download", + "local_files_only", + "revision", + "subfolder", + "token", + ) + if key in kwargs + } + processor_config_path = cached_file( + model_path, + "processor_config.json", + **cached_file_kwargs, + ) + if processor_config_path is None: + raise FileNotFoundError( + f"processor_config.json was not found for {model_path!r}." + ) + with open(processor_config_path, encoding="utf-8") as f: + vp_cfg = _cap_cfg(dict(json.load(f)["video_processor"]), is_video=True) + video_processor = Glm5NextVideoProcessor( + **{k: v for k, v in vp_cfg.items() if k != "video_processor_type"} + ) + + return cls( + image_processor=image_processor, + tokenizer=tokenizer, + video_processor=video_processor, + ) + + def __call__( + self, + images: ImageInput | None = None, + text: TextInput + | PreTokenizedInput + | list[TextInput] + | list[PreTokenizedInput] = None, + videos: VideoInput | None = None, + **kwargs: Unpack[Glm5NextProcessorKwargs], + ) -> BatchFeature: + output_kwargs = self._merge_kwargs( + Glm5NextProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + if images is not None: + image_inputs = self.image_processor( + images=images, **output_kwargs["images_kwargs"] + ) + else: + image_inputs = {} + + if videos is not None: + videos_inputs = self.video_processor( + videos=videos, **output_kwargs["videos_kwargs"] + ) + if "return_metadata" not in kwargs: + videos_inputs.pop("video_metadata") + else: + videos_inputs = {} + + if not isinstance(text, list): + text = [text] + + # Prompt updates expand the unchanged image/video markers after this call. + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop( + "return_mm_token_type_ids", False + ) + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + + if return_mm_token_type_ids: + array_ids = np.array(text_inputs["input_ids"]) + mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) + mm_token_type_ids[array_ids == self.image_token_id] = 1 + text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() + return BatchFeature( + data={**text_inputs, **image_inputs, **videos_inputs}, + tensor_type=return_tensors, + ) + + def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs): + vision_data = {} + if image_sizes is not None: + images_kwargs = Glm5NextProcessorKwargs._defaults.get("images_kwargs", {}) + images_kwargs.update(kwargs) + merge_size = ( + images_kwargs.get("merge_size", None) or self.image_processor.merge_size + ) + + num_image_patches = [ + self.image_processor.get_number_of_image_patches( + *image_size, images_kwargs + ) + for image_size in image_sizes + ] + num_image_tokens = [(n // merge_size**2) for n in num_image_patches] + vision_data.update( + { + "num_image_tokens": num_image_tokens, + "num_image_patches": num_image_patches, + } + ) + + if video_sizes is not None: + videos_kwargs = Glm5NextProcessorKwargs._defaults.get("videos_kwargs", {}) + videos_kwargs.update(kwargs) + merge_size = ( + videos_kwargs.get("merge_size", None) or self.video_processor.merge_size + ) + num_video_patches = [ + self.video_processor.get_number_of_video_patches( + *video_size, videos_kwargs + ) + for video_size in video_sizes + ] + num_video_tokens = [(n // merge_size**2) for n in num_video_patches] + vision_data["num_video_tokens"] = num_video_tokens + + return MultiModalData(**vision_data) + + def post_process_image_text_to_text( + self, + generated_outputs, + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + **kwargs, + ): + return self.tokenizer.batch_decode( + generated_outputs, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + +__all__ = [ + "Glm5NextImageProcessor", + "Glm5NextVideoProcessor", + "Glm5NextProcessor", + "smart_resize", +] diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 6b89f5c332..d501c68f2d 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -28,6 +28,10 @@ "qwen3_5_moe_text", } +# Cache-entry page sizes accepted by the paged MQA-logits kernels. Compressed +# indexer storage blocks are virtually split into one of these sizes. +PAGED_MQA_PAGE_SIZES = (32, 64) + def should_auto_disable_deep_gemm(model_type: str | None) -> bool: """Check if DeepGemm should be auto-disabled for this model on Blackwell. diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 4425fae48b..a00d0362fd 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -570,6 +570,7 @@ class AttentionMetadataBuilder(ABC, Generic[M]): # Does this backend/builder support updating the block table in existing # metadata supports_update_block_table: bool = False + requires_block_table_width: ClassVar[bool] = False @abstractmethod def __init__( diff --git a/vllm/v1/attention/backends/ddtree_branch_triton.py b/vllm/v1/attention/backends/ddtree_branch_triton.py index 21e38930bf..22a2c9f046 100644 --- a/vllm/v1/attention/backends/ddtree_branch_triton.py +++ b/vllm/v1/attention/backends/ddtree_branch_triton.py @@ -77,10 +77,7 @@ def _ddtree_paged_attention_kernel( offs_d = tl.arange(0, BLOCK_D) d_mask = offs_d < head_size q_ptrs = ( - query - + (q_start + row) * q_stride_t - + q_head * q_stride_h - + offs_d * q_stride_d + query + (q_start + row) * q_stride_t + q_head * q_stride_h + offs_d * q_stride_d ) q_vec = tl.load(q_ptrs, mask=d_mask & active, other=0.0).to(tl.float32) @@ -101,18 +98,10 @@ def _ddtree_paged_attention_kernel( for _ in range(max_q_len): cur_valid = active & (cur >= 0) & (cur < max_q_len) & (cur < q_len) visible = visible | ( - (local == cur) - & (local >= 0) - & (local < q_len) - & n_mask - & cur_valid + (local == cur) & (local >= 0) & (local < q_len) & n_mask & cur_valid ) safe_cur = tl.maximum(tl.minimum(cur, max_q_len - 1), 0) - parent_ptr = ( - parent_ids - + req_i * parent_stride0 - + safe_cur * parent_stride1 - ) + parent_ptr = parent_ids + req_i * parent_stride0 + safe_cur * parent_stride1 parent = tl.load(parent_ptr, mask=cur_valid, other=0).to(tl.int32) # In vLLM DDTree metadata, -1 means "root child". The verifier has # a synthetic prefix/root row at local slot 0, so root children must @@ -126,9 +115,7 @@ def _ddtree_paged_attention_kernel( block_offsets = offs_n - (offs_n // block_size) * block_size block_ptrs = ( - block_table - + req_i * block_stride0 - + (offs_n // block_size) * block_stride1 + block_table + req_i * block_stride0 + (offs_n // block_size) * block_stride1 ) block_ids = tl.load(block_ptrs, mask=n_mask, other=0).to(tl.int64) diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 3117c01660..c58258410f 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -17,6 +17,7 @@ if current_platform.is_cuda(): from vllm._custom_ops import reshape_and_cache_flash + try: from vllm.vllm_flash_attn import ( # type: ignore[attr-defined] flash_attn_varlen_func, diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index e6a6f15e65..cea126b91b 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -16,6 +16,7 @@ import os import time from collections.abc import Callable +from contextlib import suppress from dataclasses import dataclass from functools import partial from typing import cast @@ -23,6 +24,7 @@ import torch import vllm.envs as envs +from vllm.config.speculative import get_dflash_model_draft_tokens from vllm.forward_context import CUDAGRAPH_VARIANT_LONG_CONTEXT from vllm.logger import init_logger from vllm.platforms import current_platform @@ -1274,8 +1276,30 @@ def _get_sm70_splitd_d256_ops(): _sm70_splitd_d256_ops_checked = True try: - # Importing the interface loads the vendored FA2 torch library. - from vllm.vllm_flash_attn import flash_attn_interface # noqa: F401 + required_ops = ( + "sm70_d256_splitd_n32_dense_fwd", + "sm70_d256_splitd_n32_paged_fwd", + ) + with suppress(ImportError): + # Importing the interface loads the bundled FA2 torch library. + from vllm.vllm_flash_attn import flash_attn_interface # noqa: F401 + + namespace = getattr(torch.ops, "_vllm_fa2_C", None) + if namespace is None or not all( + hasattr(namespace, op_name) for op_name in required_ops + ): + # A partially cached Python interface can import successfully + # without registering its native operators. Source-overlay + # deployments can also intentionally keep the extension outside + # the checkout. In both cases, load only an explicitly selected + # sidecar and then validate the actual operator capability below. + library_path = os.getenv("VLLM_SM70_FA2_D256_LIBRARY") + if library_path is not None: + torch.ops.load_library(library_path) + logger.info( + "Loaded external SM70 D256 prefill library from %s.", + library_path, + ) dense = torch.ops._vllm_fa2_C.sm70_d256_splitd_n32_dense_fwd paged = torch.ops._vllm_fa2_C.sm70_d256_splitd_n32_paged_fwd @@ -1285,7 +1309,7 @@ def _get_sm70_splitd_d256_ops(): None, ) _sm70_splitd_d256_ops = (dense, paged, splitkv3) - except (AttributeError, ImportError, RuntimeError) as exc: + except (AttributeError, ImportError, OSError, RuntimeError) as exc: _sm70_splitd_d256_ops = None logger.warning_once( "SM70 D256 exact-prefill operators are unavailable (%s: %s). " @@ -1308,8 +1332,13 @@ def _get_sm70_d256_gqa_architecture_op(): _sm70_d256_gqa_architecture_op_checked = True try: - # Importing the interface loads the vendored FA2 torch library. - from vllm.vllm_flash_attn import flash_attn_interface # noqa: F401 + # The Split-D loader also resolves an explicit source-overlay + # sidecar. Calling it here keeps both operator families on one binary. + if not hasattr( + torch.ops._vllm_fa2_C, + "sm70_d256_gqa_architecture_fwd", + ): + _get_sm70_splitd_d256_ops() _sm70_d256_gqa_architecture_op = getattr( torch.ops._vllm_fa2_C, @@ -1451,7 +1480,7 @@ def _should_use_prefill_d256_gqa_architecture( softmax_scale: float, architecture_op: Callable[..., torch.Tensor] | None, ) -> bool: - """Gate the measured Q8000/KV40K..128K/Hq6/Hkv1/D256 family.""" + """Gate the stable Q8000/KV16K..256K/Hq6/Hkv1/D256 family.""" return ( envs.VLLM_FLASH_V100_PREFILL_D256_GQA_ARCH_128K_EXPERIMENTAL and architecture_op is not None @@ -1462,7 +1491,7 @@ def _should_use_prefill_d256_gqa_architecture( and value.shape == key.shape and max_seqlen_q == 8000 and max_seqlen_k == key.shape[1] - and 40000 <= max_seqlen_k <= 128000 + and 16000 <= max_seqlen_k <= 256000 and max_seqlen_k % 8000 == 0 and query.dtype == torch.float16 and key.dtype == query.dtype @@ -1632,7 +1661,7 @@ def _try_sm70_fa2_d256_prefill( if not _logged_prefill_d256_gqa_architecture: logger.info( "FLASH_ATTN_V100 SM70 D256 GQA " - "8K-by-40K..128K architecture route active." + "8K-by-16K..256K architecture route active." ) _logged_prefill_d256_gqa_architecture = True _record_route("prefill_dense_d256_gqa_arch_long") @@ -3071,7 +3100,8 @@ def __init__(self, *args, **kwargs): self._is_dflash_selector_target = bool( use_dflash and not self._is_speculative_draft_model - and getattr(spec_config, "num_speculative_tokens", None) == 7 + and getattr(spec_config, "num_speculative_tokens", None) in (7, 15) + and get_dflash_model_draft_tokens(spec_config) == 7 and selector_engine ) self._use_sm70_dflash2_fused_smallq_metadata = bool( @@ -5150,15 +5180,18 @@ def _dflash2_grouped_verify_allowed( >= self.dflash2_grouped_verify_min_model_len and getattr(attn_metadata, "causal", True) and self._flash_v100_window_size(causal=True) == (-1, -1) - and num_query_tokens == 8 - and tuple(query.shape) == (8, 6, 256) + and num_query_tokens in (8, 16) + and tuple(query.shape) == (num_query_tokens, 6, 256) and query.dtype == torch.float16 and query.is_contiguous() and key_cache.ndim == 4 and value_cache.ndim == 4 and key_cache.device == query.device and value_cache.device == query.device - and key_cache.shape[1] in (1648, 3296) + # q15 LABD increases the aligned hybrid-cache page from the + # block-8 service's 1648/3296 layout to 1728/3456. The grouped + # operator's runtime-stride implementation is exact for both. + and key_cache.shape[1] in (1648, 1728, 3296, 3456) and tuple(key_cache.shape[2:]) == (1, 256) and tuple(value_cache.shape) == tuple(key_cache.shape) and key_cache.dtype == torch.uint8 @@ -5226,7 +5259,8 @@ def _call_dflash2_grouped_verify( if not _logged_prefill_smallq_grouped_verify: logger.info( "FLASH_ATTN_V100 DFlash2 exact grouped verifier active " - "(q8/H6/Hkv1/D256, FP8 E5M2 KV, one-pass)." + "(q%d/H6/Hkv1/D256, FP8 E5M2 KV, one-pass).", + query.shape[0], ) _logged_prefill_smallq_grouped_verify = True self.flash_attn_grouped_verify_paged( diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2171f74647..0abe119e39 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -155,6 +155,60 @@ def get_kv_cache_stride_order( return (0, 1, 2) +class KpoolTailBackend(DeepseekV32IndexerBackend): + """Storage-only backend for the GLM-5.3-Flash kpool tail cache. + + The tail cache holds the in-progress pool's raw K + gate score packed into + one ``[num_blocks, 2, block_size, head_dim]`` bf16 tensor (K at index 0, + gate score at index 1) and never runs attention. It reuses the indexer + metadata builder (token-granular when ``compress_ratio == 1``) but exposes a + 4D K||score shape and unrestricted head/block sizes: the indexer backend + hard-requires ``head_size in {32,64,128}`` and ``kernel_block_size 64``, + which the tail's ``block_size == kpool`` (e.g. 4) violates. The 4D + ``[2, block_size, head_dim]`` layout keeps K and score as separate block + halves so the connectors' non-MLA K/V half-split transfers them correctly. + """ + + @staticmethod + def get_name() -> str: + return "KPOOL_TAIL" + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [] # supports any (K+score packed into head_size) + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + # block_size == kpool; no sub-block splitting (kernel_block_size == + # block_size), so the tensor stays [num_blocks, 2, kpool, head_dim]. + return [MultipleOf(1)] + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + # head_size carries K+score packed (== 2 * head_dim); split into + # [2, block_size, head_dim] so K (idx 0) and score (idx 1) are separate + # block halves. + assert num_kv_heads == 1 + assert head_size % 2 == 0, "KpoolTailSpec head_size must be 2*head_dim" + return (num_blocks, 2, block_size, head_size // 2) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + return (0, 1, 2, 3) + + @staticmethod + def get_builder_cls() -> type["KpoolTailMetadataBuilder"]: # type: ignore[override] + return KpoolTailMetadataBuilder + + class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend): @staticmethod def get_name() -> str: @@ -182,6 +236,13 @@ class DeepseekV32IndexerPrefillChunkMetadata: @dataclass class DeepseekV32IndexerPrefillMetadata: chunks: list[DeepseekV32IndexerPrefillChunkMetadata] + # Host-side max prefill seq len (== max token position + 1 across the + # step's prefill tokens; exact for prefill rows per the assert in build()). + # Lets the indexer layer's short-sequence check branch without the + # device sync a positions.max().item() would cost on every layer. + # -1 = unknown (other construction sites): the layer falls back to the + # device-side check. + max_prefill_seq_len: int = -1 @dataclass @@ -195,6 +256,9 @@ class DeepSeekV32IndexerDecodeMetadata: decode_lens: torch.Tensor requires_padding: bool schedule_metadata: torch.Tensor + per_req_decode_lens: torch.Tensor | None = None + decode_is_uniform: bool = True + write_max_decode_len: int = 0 @dataclass @@ -221,6 +285,114 @@ class DeepseekV32IndexerMetadata: prefill: DeepseekV32IndexerPrefillMetadata | None = None +def compute_kpool_tail_slot_mapping( + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + query_start_loc: torch.Tensor, + positions: torch.Tensor, + num_actual_tokens: int, + num_reqs: int, + kpool: int, +) -> torch.Tensor: + """Circular tail slots: every token of request r lands in r's own block. + + The generic per-group slot kernel maps ``pos -> bt[req][pos // bs] * bs + + pos % bs`` (``_compute_slot_mappings_kernel``). The tail group's block + table only ever has its FIRST column written -- ``KpoolTailManager`` + allocates exactly one block per request and never grows -- so for any + token at ``pos >= kpool`` that kernel reads a zero column and the slot + collapses onto physical tail block 0. All concurrently running requests + then share one ``kpool``-slot ring and corrupt each other's pool + compression (seed / stash / completion all target block 0). This maps + ``slot = own_block * kpool + pos % kpool`` instead, which is the layout + the tail kernels and ``KpoolTailManager`` are designed around. + + Pure torch (no Triton, no device sync): the indexer op consumes the tail + slot mapping on its eager break, so the returned tensor need not be the + persistent ``BlockTables`` buffer. + """ + out = slot_mapping.clone() + if num_actual_tokens == 0: + return out + device = slot_mapping.device + tokens = torch.arange(num_actual_tokens, device=device) + # searchsorted(right=True): token i in [qsl[r], qsl[r+1]) -> request r. + req = torch.searchsorted(query_start_loc, tokens, right=True) - 1 + req = req.clamp_(min=0, max=num_reqs - 1) + own_block = block_table[:num_reqs, 0].index_select(0, req).to(torch.int64) + pos = positions[:num_actual_tokens].to(torch.int64) + out[:num_actual_tokens] = own_block * kpool + torch.remainder(pos, kpool) + return out + + +class KpoolTailMetadataBuilder(AttentionMetadataBuilder): + """Lean metadata builder for the kpool tail cache. + + The tail is storage-only (no attention / MQA-logits), so it skips the + DeepseekV32 indexer builder's DeepGEMM paged-MQA path -- which requires + ``block_kv in {32,64}`` and asserts on the tail's ``block_size == kpool``. + It exports only the group's token-granular ``slot_mapping`` plus the + prefill/decode counts; the indexer op reads + ``attn_metadata[tail_prefix].slot_mapping``. The slot mapping is + recomputed with the circular per-request layout (see + :func:`compute_kpool_tail_slot_mapping`) instead of reusing the generic + per-group kernel output, which cannot express a 1-block-per-request + circular buffer. + """ + + _cudagraph_support = AttentionCGSupport.ALWAYS + supports_update_block_table = False + reorder_batch_threshold = None + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ): + # No indexer-builder buffers (expanded_block_table / scheduler_metadata / + # compressed_slot_mapping) -- the tail is storage-only and exports only + # slot_mapping, which is rebuilt per step from the group's block table. + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> DeepseekV32IndexerMetadata: + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills(common_attn_metadata) + ) + slot_mapping = common_attn_metadata.slot_mapping + positions = common_attn_metadata.positions + if positions is not None: + # Circular per-request layout; the generic kernel output collapses + # onto tail block 0 for pos >= kpool (see compute_... docstring). + slot_mapping = compute_kpool_tail_slot_mapping( + slot_mapping, + common_attn_metadata.block_table_tensor, + common_attn_metadata.query_start_loc, + positions, + common_attn_metadata.num_actual_tokens, + common_attn_metadata.num_reqs, + self.kv_cache_spec.block_size, + ) + return DeepseekV32IndexerMetadata( + seq_lens=common_attn_metadata.seq_lens, + max_seq_len=common_attn_metadata.max_seq_len, + compressed_max_seq_len=common_attn_metadata.max_seq_len, + slot_mapping=slot_mapping, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + prefill=None, + decode=None, + ) + + def get_max_prefill_buffer_size(vllm_config: VllmConfig): max_model_len = vllm_config.model_config.max_model_len # NOTE(Chen): 40 is a magic number for controlling the prefill buffer size. @@ -235,6 +407,7 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig): class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): + requires_block_table_width = True reorder_batch_threshold: int = 1 natively_supported_next_n_fp4: list[int] = [1, 2] # TODO (matt): integrate kernel with next_n = 4 support @@ -247,7 +420,7 @@ def get_cudagraph_support( ) -> AttentionCGSupport: return AttentionCGSupport.UNIFORM_BATCH - def __init__(self, *args, **kwargs): + def __init__(self, *args, block_table_width: int | None = None, **kwargs): super().__init__(*args, **kwargs) scheduler_config = self.vllm_config.scheduler_config # NOTE(Chen):an estimated max size of flattened_kv. Need to double check. @@ -293,6 +466,15 @@ def __init__(self, *args, **kwargs): dtype=torch.int32, device=self.device, ) + # Persistent mirror of the original per-request decode_lens (length + # num_decodes), captured in build() before _prepare_decode_tensors + # flattens decode_lens_buffer to all-1s. Persistent so FULL-cudagraph + # replay reads a stable pointer (a per-step clone would move address). + self.per_req_decode_lens_buffer = torch.zeros( + (scheduler_config.max_num_batched_tokens,), + dtype=torch.int32, + device=self.device, + ) # Shared workspace for decode seq_lens. Native MTP views this as # (B, max_decode_len) at runtime, keeping context_lens contiguous even # when max_decode_len is smaller than next_n. @@ -309,14 +491,15 @@ def __init__(self, *args, **kwargs): dtype=torch.int32, device=self.device, ) - max_num_blocks_per_req = cdiv( - self.vllm_config.model_config.max_model_len, - self.kv_cache_spec.block_size * get_total_cp_world_size(), - ) + if block_table_width is None: + block_table_width = cdiv( + self.vllm_config.model_config.max_model_len, + self.kv_cache_spec.block_size * get_total_cp_world_size(), + ) self.expanded_block_table_buffer = torch.zeros( ( scheduler_config.max_num_batched_tokens, - max_num_blocks_per_req, + block_table_width, ), dtype=torch.int32, device=self.device, @@ -348,6 +531,10 @@ def __init__(self, *args, **kwargs): dtype=torch.int32, device=self.device, ) + # FULL CUDA graphs require the translated decode table at a stable + # address; allocate it lazily at the fixed compressed width. + self.indexer_decode_block_table_buffer: torch.Tensor | None = None + self._max_num_batched_tokens = scheduler_config.max_num_batched_tokens def _prepare_decode_tensors( self, @@ -496,12 +683,27 @@ def build( compressed_slot_mapping = slot_mapping compressed_seq_lens = seq_lens + # Default to the raw shared table; translated below when the indexer is + # co-located with a virtually-split MLA (compress_ratio > 1). This + # translated table is the ONE block_table the indexer must use for every + # cache access (write via slot_mapping AND the top-k reads), so that the + # pool keys written at indexer-page granularity are read back from the + # same physical pages. + indexer_block_table = block_table if self.compress_ratio > 1: + kbs = getattr(self, "kernel_block_size", None) + if ( + kbs is not None + and self.kv_cache_spec.block_size != kbs + and self.kv_cache_spec.block_size % kbs == 0 + ): + factor = self.kv_cache_spec.block_size // kbs + indexer_block_table = (block_table[:, ::factor] // factor).contiguous() compressed_slot_mapping = get_compressed_slot_mapping( num_tokens, query_start_loc, seq_lens, - block_table, + indexer_block_table, self.kv_cache_spec.storage_block_size, self.compress_ratio, out=self.compressed_slot_mapping_buffer, @@ -546,7 +748,7 @@ def build( seq_lens, compressed_seq_lens, compressed_seq_lens_cpu, - common_attn_metadata.block_table_tensor, + indexer_block_table, self.compress_ratio, query_slice=query_slice, skip_kv_gather=query_slice.start > 0, @@ -554,7 +756,19 @@ def build( # Skip when total_seq_lens is 0 (i.e., no compressed token). if metadata is not None: chunks.append(metadata) - prefill_metadata = DeepseekV32IndexerPrefillMetadata(chunks) + prefill_metadata = DeepseekV32IndexerPrefillMetadata( + chunks, + max_prefill_seq_len=( + # Uncompressed (token-granular) to match `positions`. + # seq_lens_cpu is exact for the prefill rows here, and the + # step's prefill tokens always include each request's last + # token (position seq_len-1), so this equals + # positions[prefill_slice].max() + 1. + int(seq_lens_cpu[num_decodes:].max().item()) + if num_prefills > 0 + else 0 + ), + ) decode_metadata = None if num_decodes > 0: @@ -563,6 +777,12 @@ def build( out=self.decode_lens_buffer[:num_decodes], ) decode_lens = self.decode_lens_buffer[:num_decodes] + # Stash the original per-request decode_lens before + # _prepare_decode_tensors rewrites decode_lens_buffer (the flatten + # path sets every entry to 1 for the logits read). The kpool + # decode-write path uses these to scatter tokens by request on a + # variable MTP-verify batch. + self.per_req_decode_lens_buffer[:num_decodes].copy_(decode_lens) decode_lens_cpu = torch.diff( common_attn_metadata.query_start_loc_cpu[: num_decodes + 1] ) @@ -571,6 +791,14 @@ def build( block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...] max_decode_len = int(decode_lens_cpu.max().item()) + min_decode_len = int(decode_lens_cpu.min().item()) + # Host-side uniformity flag for the kpool decode-write path: the + # flatten path below rewrites decode_lens to all-1s and reports + # requires_padding=False even for a variable MTP-verify batch, so + # the write path can't reuse requires_padding. Compute it here (on + # the CPU copy, so no runtime .item() that would break cudagraph + # capture) and hand down max_decode_len as the scatter lmax. + write_is_uniform = min_decode_len == max_decode_len next_n = 1 + self.num_speculative_tokens use_native = not self.use_flattening and max_decode_len <= next_n @@ -589,6 +817,27 @@ def build( ) ) + if self.compress_ratio > 1: + kbs = getattr(self, "kernel_block_size", None) + if ( + kbs is not None + and self.kv_cache_spec.block_size != kbs + and self.kv_cache_spec.block_size % kbs == 0 + ): + factor = self.kv_cache_spec.block_size // kbs + compressed = block_table[:, ::factor] // factor + rows, cols = compressed.shape + if self.indexer_decode_block_table_buffer is None: + self.indexer_decode_block_table_buffer = torch.zeros( + (self._max_num_batched_tokens, cols), + dtype=torch.int32, + device=self.device, + ) + self.indexer_decode_block_table_buffer[:rows, :cols].copy_( + compressed + ) + block_table = self.indexer_decode_block_table_buffer[:rows, :cols] + # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores # compressed tokens. Convert uncompressed seq_lens to compressed. if self.compress_ratio > 1: @@ -627,6 +876,9 @@ def build( decode_lens=decode_lens, requires_padding=requires_padding, schedule_metadata=self.scheduler_metadata_buffer, + per_req_decode_lens=self.per_req_decode_lens_buffer[:num_decodes], + decode_is_uniform=write_is_uniform, + write_max_decode_len=max_decode_len, ) attn_metadata = DeepseekV32IndexerMetadata( diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index a96c0dd115..e077b3f9d0 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -42,12 +42,8 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): """ FLASH_ATTN = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" - FLASH_ATTN_V100 = ( - "vllm.v1.attention.backends.flash_attn_v100.FlashAttnV100Backend" - ) - FLASHINFER_SM70 = ( - "vllm.v1.attention.backends.flashinfer_sm70.FlashInferSM70Backend" - ) + FLASH_ATTN_V100 = "vllm.v1.attention.backends.flash_attn_v100.FlashAttnV100Backend" + FLASHINFER_SM70 = "vllm.v1.attention.backends.flashinfer_sm70.FlashInferSM70Backend" FLASH_ATTN_DIFFKV = ( "vllm.v1.attention.backends.flash_attn_diffkv.FlashAttentionDiffKVBackend" ) @@ -82,6 +78,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): FLASHMLA_SPARSE = ( "vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend" ) + GLM5_SM70_SPARSE = "vllm.models.glm5next.sm70.sparse.Glm5NextSM70SparseBackend" FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend" NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend" FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" diff --git a/vllm/v1/attention/backends/short_conv_attn.py b/vllm/v1/attention/backends/short_conv_attn.py index 9c85ec5efb..efbbd1e075 100644 --- a/vllm/v1/attention/backends/short_conv_attn.py +++ b/vllm/v1/attention/backends/short_conv_attn.py @@ -1,12 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, replace +from typing import Any -from vllm.v1.attention.backend import AttentionBackend +import torch + +from vllm.config import VllmConfig +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + CommonAttentionMetadata, +) from vllm.v1.attention.backends.mamba_attn import ( BaseMambaAttentionMetadata, BaseMambaAttentionMetadataBuilder, ) +from vllm.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + compute_causal_conv1d_metadata, + mamba_get_block_table_tensor, +) +from vllm.v1.kv_cache_interface import AttentionSpec class ShortConvAttentionBackend(AttentionBackend): @@ -32,3 +47,515 @@ class ShortConvAttentionMetadataBuilder( BaseMambaAttentionMetadataBuilder[ShortConvAttentionMetadata] ): metadata_cls = ShortConvAttentionMetadata + + +@dataclass +class PleShortConvAttentionMetadata(ShortConvAttentionMetadata): + # Number of speculative-decode (multi-query / MTP) requests and the total + # number of tokens they contribute. These are 0 when spec-decode is off. + num_spec_decodes: int = 0 + num_spec_decode_tokens: int = 0 + num_actual_tokens: int = 0 + + # Max query length among spec-decode requests (== num_speculative_tokens + 1). + # Used as ``max_query_len`` for the varlen spec causal_conv1d_update. + spec_query_len: int = 1 + + # Max query length among the non-spec *prefill* requests, precomputed + # CPU-side in the builder. The dilated PLE short-conv uses it to size its + # packing buffer without a device->host sync (``lengths.max().item()``). + # 0 when there are no prefill requests. + max_prefill_query_len: int = 0 + query_start_loc: torch.Tensor | None = None + + # ``state_indices_tensor`` keeps the historical (non-spec) layout used by + # all existing short-conv consumers: the conv-state slot for each regular + # decode followed by each prefill request. When spec-decode is active this + # only covers the non-spec requests. + state_indices_tensor: torch.Tensor | None = None + has_initial_states_d: torch.Tensor | None = None + + # ``non_spec_query_start_loc`` is the varlen cumulative token offset over + # the non-spec requests only (decodes then prefills). It equals + # ``query_start_loc`` when there are no spec-decode requests. + non_spec_query_start_loc: torch.Tensor | None = None + + # Speculative-decode (MTP) conv metadata. Only column 0 of the block table + # is needed for the convolution state, so these tensors are 1-D over the + # spec-decode requests. + spec_query_start_loc: torch.Tensor | None = None # [num_spec_decodes + 1] + spec_state_indices_tensor: torch.Tensor | None = None # [num_spec_decodes] + spec_sequence_masks: torch.Tensor | None = None # [batch] + spec_token_indx: torch.Tensor | None = None + non_spec_token_indx: torch.Tensor | None = None + num_decode_draft_tokens_cpu: torch.Tensor | None = None + + +class PleShortConvAttentionBackend(ShortConvAttentionBackend): + @staticmethod + def get_name() -> str: + return "PLE_SHORT_CONV_ATTN" + + @staticmethod + def get_builder_cls() -> type["PleShortConvAttentionMetadataBuilder"]: + return PleShortConvAttentionMetadataBuilder + + +class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): + metadata_cls = PleShortConvAttentionMetadata + # Spec-decode requires a uniform (multi-token) decode batch for full + # CUDA graph capture, matching the GDN backend. + _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH + reorder_batch_threshold: int = 1 + supports_update_block_table = False + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.num_spec = self.num_spec_tokens + self.use_full_cuda_graph = ( + self.compilation_config.cudagraph_mode.has_full_cudagraphs() + ) + + max_num_seqs = vllm_config.scheduler_config.max_num_seqs + max_capture_size = self.compilation_config.max_cudagraph_capture_size + self.decode_cudagraph_max_bs = max_num_seqs + self.decode_cudagraph_max_tokens = max_num_seqs * (self.num_spec + 1) + if max_capture_size is not None: + self.decode_cudagraph_max_bs = min( + self.decode_cudagraph_max_bs, max_capture_size + ) + self.decode_cudagraph_max_tokens = min( + self.decode_cudagraph_max_tokens, max_capture_size + ) + + # Persistent buffers reused during full CUDA graph capture and replay. + self.spec_state_indices_tensor = torch.empty( + (self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device + ) + self.spec_sequence_masks = torch.empty( + (self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device + ) + self.spec_token_indx = torch.empty( + (self.decode_cudagraph_max_tokens,), dtype=torch.int32, device=device + ) + self.non_spec_token_indx = torch.empty( + (self.decode_cudagraph_max_tokens,), dtype=torch.int32, device=device + ) + self.spec_query_start_loc = torch.empty( + (self.decode_cudagraph_max_bs + 1,), dtype=torch.int32, device=device + ) + self.num_accepted_tokens = torch.empty( + (self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device + ) + self.has_initial_states_d = torch.empty( + (self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device + ) + + def _build_non_spec_metadata( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool, + num_decode_draft_tokens_cpu: torch.Tensor | None, + **kwargs: Any, + ) -> PleShortConvAttentionMetadata: + metadata = super().build( + common_prefix_len, + common_attn_metadata, + fast_build, + num_accepted_tokens=None, + **kwargs, + ) + assert isinstance(metadata, PleShortConvAttentionMetadata) + + state_indices_d = metadata.state_indices_tensor_d + if state_indices_d is not None and state_indices_d.dim() > 1: + state_indices_d = state_indices_d[:, 0] + state_indices_p = metadata.state_indices_tensor_p + if metadata.num_prefills == 0: + assert state_indices_d is not None + # BaseMambaAttentionMetadataBuilder pads decode state indices into + # a persistent tensor for full CUDA graphs. Keep those rows so the + # PLE decode receives one cache slot per graph-padded token. + state_indices_tensor = state_indices_d + elif metadata.num_decodes == 0: + assert state_indices_p is not None + state_indices_tensor = state_indices_p[: metadata.num_prefills] + else: + assert state_indices_d is not None + assert state_indices_p is not None + state_indices_tensor = torch.cat( + (state_indices_d, state_indices_p[: metadata.num_prefills]) + ) + + has_initial_states_d = None + if metadata.num_decodes > 0: + num_computed_tokens = common_attn_metadata.compute_num_computed_tokens() + has_initial_states_d = num_computed_tokens[: metadata.num_decodes] > 0 + if ( + self.use_full_cuda_graph + and metadata.num_prefills == 0 + and metadata.num_decodes <= self.decode_cudagraph_max_bs + ): + assert state_indices_d is not None + # Prepare tensors for CUDA graph replay. Padded rows have no + # initial state and use NULL_BLOCK_ID in state_indices_d. + num_decode_rows = state_indices_d.numel() + self.has_initial_states_d[: metadata.num_decodes].copy_( + has_initial_states_d, non_blocking=True + ) + self.has_initial_states_d[metadata.num_decodes : num_decode_rows].fill_( + False + ) + has_initial_states_d = self.has_initial_states_d[:num_decode_rows] + + max_prefill_query_len = 0 + if metadata.num_prefills > 0: + query_lens_cpu = torch.diff(common_attn_metadata.query_start_loc_cpu) + max_prefill_query_len = int( + query_lens_cpu[ + metadata.num_decodes : ( + metadata.num_decodes + metadata.num_prefills + ) + ] + .max() + .item() + ) + + return replace( + metadata, + num_actual_tokens=common_attn_metadata.num_actual_tokens, + spec_query_len=self.num_spec + 1, + max_prefill_query_len=max_prefill_query_len, + query_start_loc=common_attn_metadata.query_start_loc, + state_indices_tensor=state_indices_tensor, + has_initial_states_d=has_initial_states_d, + non_spec_query_start_loc=common_attn_metadata.query_start_loc, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) + + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + *, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + **kwargs: Any, + ) -> PleShortConvAttentionMetadata: + m = common_attn_metadata + spec_sequence_masks_cpu: torch.Tensor | None = None + # Detect speculative-decode requests. We use -1 to mark prefill and + # plain-decode requests, so any value >= 0 is a (multi-query) + # spec-decode request. + if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: + candidate_mask = num_decode_draft_tokens_cpu[: m.num_reqs] >= 0 + if bool(candidate_mask.any().item()): + spec_sequence_masks_cpu = candidate_mask + + if spec_sequence_masks_cpu is None: + return self._build_non_spec_metadata( + common_prefix_len, + common_attn_metadata, + fast_build, + num_decode_draft_tokens_cpu, + **kwargs, + ) + + del common_prefix_len, fast_build, kwargs + query_start_loc = m.query_start_loc + query_start_loc_cpu = m.query_start_loc_cpu + query_lens_cpu = torch.diff(query_start_loc_cpu) + block_table_tensor = mamba_get_block_table_tensor( + m.block_table_tensor, + m.seq_lens, + self.kv_cache_spec, + self.vllm_config.cache_config.mamba_cache_mode, + ) + + if query_start_loc.device.type == "cpu": + spec_sequence_masks = spec_sequence_masks_cpu + else: + spec_sequence_masks = async_tensor_h2d( + spec_sequence_masks_cpu.tolist(), + dtype=torch.bool, + device=query_start_loc.device, + ) + + # For causal_conv1d (non-spec prefill Triton kernel metadata). + nums_dict = None + batch_ptr = None + token_chunk_offset_ptr = None + has_initial_states_p = None + has_initial_states_d = None + num_computed_tokens_p = None + # Original request indices of the non-spec requests, ordered + # [decodes, prefills]. Used to gather per-request data consistently. + non_spec_req_idx_cpu: torch.Tensor | None = None + + query_lens = torch.diff(query_start_loc) + # Per-request classification by mask, NOT by position. With + # spec-decode, the front decode group can contain both spec-decode + # requests and plain non-spec single-token decodes. Spec requests are + # therefore not guaranteed to occupy the first num_spec_decodes slots. + non_spec_mask_cpu = ~spec_sequence_masks_cpu + decode_mask_cpu = non_spec_mask_cpu & (query_lens_cpu == 1) + prefill_mask_cpu = non_spec_mask_cpu & (query_lens_cpu > 1) + + num_spec_decodes = int(spec_sequence_masks_cpu.sum().item()) + num_decodes = int(decode_mask_cpu.sum().item()) + num_prefills = int(prefill_mask_cpu.sum().item()) + num_decode_tokens = num_decodes + num_prefill_tokens = int(query_lens_cpu[prefill_mask_cpu].sum().item()) + num_spec_decode_tokens = int( + query_lens_cpu[spec_sequence_masks_cpu].sum().item() + ) + # Max prefill query length, CPU-side (no device-to-host sync) for the + # PLE dilated short-conv packing buffer. + max_prefill_query_len = ( + int(query_lens_cpu[prefill_mask_cpu].max().item()) + if num_prefills > 0 + else 0 + ) + + # Original request indices grouped as + # [spec | non-spec decode | non-spec prefill]; each group keeps the + # original (already reordered) relative order via a stable nonzero. + spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0] + decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0] + prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0] + non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu)) + spec_req_idx = spec_req_idx_cpu.to(query_start_loc.device) + non_spec_req_idx = non_spec_req_idx_cpu.to(query_start_loc.device) + + if num_decodes == 0 and num_prefills == 0: + # Pure speculative-decode batch: all real tokens are spec tokens. + spec_token_indx = torch.arange( + num_spec_decode_tokens, + dtype=torch.int32, + device=query_start_loc.device, + ) + non_spec_token_indx = torch.empty( + 0, dtype=torch.int32, device=query_start_loc.device + ) + spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0] + non_spec_state_indices_tensor = None + spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] + non_spec_query_start_loc = None + non_spec_query_start_loc_cpu = None + else: + # Mixed batch: build a per-token group key consistent with the + # request grouping above (spec=0 | decode=1 | prefill=2) and a + # stable sort, so tokens of each request stay contiguous and in + # request order. This yields spec tokens first, then the non-spec + # [decode, prefill] tokens. + req_group = torch.full( + (m.num_reqs,), + 2, + dtype=torch.int64, + device=query_start_loc.device, + ) + req_group[spec_req_idx] = 0 + req_group[decode_req_idx_cpu.to(query_start_loc.device)] = 1 + token_group = torch.repeat_interleave(req_group, query_lens) + token_perm = torch.argsort(token_group, stable=True) + spec_token_indx = token_perm[:num_spec_decode_tokens] + non_spec_token_indx = token_perm[num_spec_decode_tokens:] + + spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0] + non_spec_state_indices_tensor = block_table_tensor[non_spec_req_idx, 0] + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum(query_lens[spec_req_idx], dim=0, out=spec_query_start_loc[1:]) + non_spec_query_start_loc = torch.zeros( + num_decodes + num_prefills + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[non_spec_req_idx], + dim=0, + out=non_spec_query_start_loc[1:], + ) + non_spec_query_start_loc_cpu = torch.zeros( + num_decodes + num_prefills + 1, dtype=torch.int32 + ) + torch.cumsum( + query_lens_cpu[non_spec_req_idx_cpu], + dim=0, + out=non_spec_query_start_loc_cpu[1:], + ) + + assert num_accepted_tokens is not None + # Accepted-token counts must follow the same request order as the + # speculative state indices. + num_accepted_tokens = num_accepted_tokens[ + spec_req_idx_cpu.to(num_accepted_tokens.device) + ] + + # Compute the conv-state slots for the non-spec decode/prefill split, + # plus the initial-state masks and Triton causal_conv1d metadata. + if non_spec_state_indices_tensor is None: + state_indices_tensor = block_table_tensor[:0, 0] + else: + state_indices_tensor = non_spec_state_indices_tensor + + # Build the regular decode/prefill state metadata inherited from the + # generic short-conv metadata contract. + query_start_loc_p = None + query_start_loc_d = None + state_indices_tensor_p = None + state_indices_tensor_d = None + if num_decodes > 0 or num_prefills > 0: + num_computed_tokens = m.compute_num_computed_tokens() + if non_spec_req_idx_cpu is not None: + non_spec_req_idx = non_spec_req_idx_cpu.to(num_computed_tokens.device) + num_computed_tokens = num_computed_tokens[non_spec_req_idx] + + state_indices_tensor_d = state_indices_tensor[:num_decodes] + state_indices_tensor_p = state_indices_tensor[ + num_decodes : num_decodes + num_prefills + ] + if num_decodes > 0: + has_initial_states_d = num_computed_tokens[:num_decodes] > 0 + assert non_spec_query_start_loc is not None + query_start_loc_d = non_spec_query_start_loc[: num_decodes + 1] + if num_prefills > 0: + num_computed_tokens_p = num_computed_tokens[ + num_decodes : num_decodes + num_prefills + ] + has_initial_states_p = num_computed_tokens_p > 0 + assert non_spec_query_start_loc is not None + assert non_spec_query_start_loc_cpu is not None + query_start_loc_p = ( + non_spec_query_start_loc[num_decodes:] - num_decode_tokens + ) + query_start_loc_p_cpu = ( + non_spec_query_start_loc_cpu[num_decodes:] - num_decode_tokens + ) + if query_start_loc.device.type != "cpu": + nums_dict, batch_ptr, token_chunk_offset_ptr = ( + compute_causal_conv1d_metadata( + query_start_loc_p_cpu, + device=query_start_loc.device, + ) + ) + + # Prepare persistent tensors for CUDA graph capture and replay. + # ``m.num_actual_tokens`` is already padded by the model runner. + # Request-level buffers use ``m.num_reqs`` while token-level buffers + # use their independently bounded token count. + batch_size = m.num_reqs + if ( + self.use_full_cuda_graph + and num_prefills == 0 + and num_decodes == 0 + and spec_sequence_masks is not None + and num_spec_decodes <= self.decode_cudagraph_max_bs + and num_spec_decode_tokens <= self.decode_cudagraph_max_tokens + ): + assert spec_state_indices_tensor is not None + self.spec_state_indices_tensor[:num_spec_decodes].copy_( + spec_state_indices_tensor, non_blocking=True + ) + spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size] + spec_state_indices_tensor[num_spec_decodes:].fill_(NULL_BLOCK_ID) + + self.spec_sequence_masks[:batch_size].copy_( + spec_sequence_masks[:batch_size], non_blocking=True + ) + spec_sequence_masks = self.spec_sequence_masks[:batch_size] + + assert spec_query_start_loc is not None + self.spec_query_start_loc[: num_spec_decodes + 1].copy_( + spec_query_start_loc, non_blocking=True + ) + spec_num_query_tokens = spec_query_start_loc[-1] + spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1] + spec_query_start_loc[num_spec_decodes + 1 :].fill_(spec_num_query_tokens) + + assert num_accepted_tokens is not None + self.num_accepted_tokens[:num_spec_decodes].copy_( + num_accepted_tokens, non_blocking=True + ) + num_accepted_tokens = self.num_accepted_tokens[:batch_size] + num_accepted_tokens[num_spec_decodes:].fill_(1) + + return PleShortConvAttentionMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_reqs=m.num_reqs, + num_spec_decodes=num_spec_decodes, + num_spec_decode_tokens=num_spec_decode_tokens, + num_actual_tokens=m.num_actual_tokens, + spec_query_len=self.num_spec + 1, + max_prefill_query_len=max_prefill_query_len, + query_start_loc=query_start_loc, + state_indices_tensor=state_indices_tensor, + has_initial_states_p=has_initial_states_p, + has_initial_states_d=has_initial_states_d, + non_spec_query_start_loc=non_spec_query_start_loc, + spec_query_start_loc=spec_query_start_loc, + spec_state_indices_tensor=spec_state_indices_tensor, + spec_sequence_masks=spec_sequence_masks, + spec_token_indx=spec_token_indx, + non_spec_token_indx=non_spec_token_indx, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + nums_dict=nums_dict, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=token_chunk_offset_ptr, + query_start_loc_p=query_start_loc_p, + query_start_loc_d=query_start_loc_d, + state_indices_tensor_p=state_indices_tensor_p, + state_indices_tensor_d=state_indices_tensor_d, + num_computed_tokens_p=num_computed_tokens_p, + block_idx_last_scheduled_token=None, + block_idx_first_scheduled_token_p=None, + block_idx_last_computed_token=None, + block_idx_last_scheduled_token_prev_step=None, + seq_lens=m.seq_lens, + ) + + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ) -> PleShortConvAttentionMetadata: + """Build metadata for full CUDA graph capture. + + Currently, only decode is supported for full CUDA graphs with + short-conv. + """ + m = common_attn_metadata + assert ( + m.num_reqs <= self.decode_cudagraph_max_bs + and m.num_actual_tokens <= self.decode_cudagraph_max_tokens + ), ( + "ShortConv only supports decode-only full CUDAGraph capture. " + f"Make sure batch size ({m.num_reqs}) <= " + f"cudagraph capture size ({self.decode_cudagraph_max_bs}) and " + f"number of tokens ({m.num_actual_tokens}) <= " + f"token capture size ({self.decode_cudagraph_max_tokens})." + ) + + if self.use_spec_decode: + num_accepted_tokens = torch.diff(m.query_start_loc) + num_decode_draft_tokens_cpu = (num_accepted_tokens - 1).cpu() + return self.build( + 0, + m, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) + return self.build(0, m) diff --git a/vllm/v1/attention/ops/triton_prefill_attention.py b/vllm/v1/attention/ops/triton_prefill_attention.py index 4e28e95e95..8450ba1eda 100644 --- a/vllm/v1/attention/ops/triton_prefill_attention.py +++ b/vllm/v1/attention/ops/triton_prefill_attention.py @@ -184,11 +184,7 @@ def get_block_size(dtype: torch.dtype, head_dim: int | None = None) -> int: 80 ): return 128 - elif ( - current_platform.is_cuda_alike() - and head_dim is not None - and head_dim >= 256 - ): + elif current_platform.is_cuda_alike() and head_dim is not None and head_dim >= 256: # SM70/SM75 only expose 96KB shared memory per block. With head_dim=256, # BLOCK=64 requires 128KB in this Triton kernel. return 32 diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index ef5b256eb6..0d59a03535 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -558,9 +558,7 @@ def kernel_unified_attention( if PV_INPUT_PRECISION == "": acc += tl.dot(P.to(V.dtype), V) else: - acc += tl.dot( - P.to(V.dtype), V, input_precision=PV_INPUT_PRECISION - ) + acc += tl.dot(P.to(V.dtype), V, input_precision=PV_INPUT_PRECISION) # ---- Epilogue --------------------------------------------------------- if IS_3D: @@ -920,12 +918,8 @@ def unified_attention( sm70_prefill_tile_size = envs.VLLM_SM70_TRITON_ATTN_PREFILL_TILE_SIZE sm70_decode_tile_size = envs.VLLM_SM70_TRITON_ATTN_DECODE_TILE_SIZE sm70_num_warps = envs.VLLM_SM70_TRITON_ATTN_NUM_WARPS - sm70_prefill_num_warps = ( - envs.VLLM_SM70_TRITON_ATTN_PREFILL_NUM_WARPS - ) - sm70_decode_num_warps = ( - envs.VLLM_SM70_TRITON_ATTN_DECODE_NUM_WARPS - ) + sm70_prefill_num_warps = envs.VLLM_SM70_TRITON_ATTN_PREFILL_NUM_WARPS + sm70_decode_num_warps = envs.VLLM_SM70_TRITON_ATTN_DECODE_NUM_WARPS sm70_safe_defaults = envs.VLLM_SM70_TRITON_ATTN_SAFE_DEFAULTS sm70_qk_input_precision = _validate_sm70_triton_attn_input_precision( envs.VLLM_SM70_TRITON_ATTN_QK_INPUT_PRECISION, diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 224b4d98d7..81f904b1c7 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -433,9 +433,14 @@ def __init__( # different KV cache groups have different block sizes, the actual block size # can be a multiple of hash_block_size. self.hash_block_size = hash_block_size + participating_block_sizes = [ + group.kv_cache_spec.block_size + for group in kv_cache_config.kv_cache_groups + if group.kv_cache_spec.prefix_cacheable + ] assert all( - g.kv_cache_spec.block_size % hash_block_size == 0 - for g in kv_cache_config.kv_cache_groups + block_size % hash_block_size == 0 + for block_size in participating_block_sizes ), "block_size must be divisible by hash_block_size" assert dcp_world_size == 1, "DCP not support hybrid attn now." assert pcp_world_size == 1, "PCP not support hybrid attn now." @@ -451,6 +456,13 @@ def verify_and_split_kv_cache_groups(self) -> None: ] = [] for i, g in enumerate(self.kv_cache_config.kv_cache_groups): + # Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash + # kpool tail): their blocks are per-request scratch, never + # shareable, so they must not participate in hit lookup (their + # manager-level hooks already no-op). Their slot in the per-group + # hit tuple stays empty. + if not g.kv_cache_spec.prefix_cacheable: + continue manager_cls = self.single_type_managers[i].__class__ spec = g.kv_cache_spec @@ -465,8 +477,8 @@ def verify_and_split_kv_cache_groups(self) -> None: else: attention_groups.append((spec, [i], manager_cls)) - assert len(attention_groups) > 1, ( - "HybridKVCacheCoordinator requires at least two attention groups." + assert attention_groups, ( + "HybridKVCacheCoordinator requires at least one cacheable group." ) # Put full attention first: its efficient left-to-right scan provides diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 9f1f447340..6ebdcf8b78 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -7,10 +7,10 @@ import math import os from collections import defaultdict -from collections.abc import Callable, Iterable, Iterator, Sequence +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass, replace from functools import partial -from typing import Any, NewType, TypeAlias, cast, overload +from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload from vllm import envs from vllm.config import VllmConfig @@ -21,8 +21,10 @@ from vllm.utils.torch_utils import get_dtype_size from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, + CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, + KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, @@ -880,18 +882,14 @@ def get_max_concurrency_for_kv_cache_config( """ Get the maximum concurrency for the given KV cache configuration. """ - num_layer_per_group = max( - len(group.layer_names) for group in kv_cache_config.kv_cache_groups - ) - max_memory_usage_per_request = num_layer_per_group * max_memory_usage_bytes( - vllm_config, (group.kv_cache_spec for group in kv_cache_config.kv_cache_groups) - ) - memory_per_block = ( - kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes - * num_layer_per_group + num_blocks_per_request = sum( + cdiv( + group.kv_cache_spec.max_memory_usage_bytes(vllm_config), + group.kv_cache_spec.page_size_bytes, + ) + for group in kv_cache_config.kv_cache_groups ) - num_block_per_request = cdiv(max_memory_usage_per_request, memory_per_block) - max_concurrency = kv_cache_config.num_blocks / num_block_per_request + max_concurrency = kv_cache_config.num_blocks / num_blocks_per_request return max_concurrency @@ -916,6 +914,11 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs ): return kv_cache_groups[0].kv_cache_spec.page_size_bytes + if (glm5 := _glm5_next_tensor_layout(kv_cache_groups)) is not None: + _, _, mla_names, idx_names, mla_page, idx_page, _, _ = glm5 + return len(mla_names) * mla_page + len(idx_names) * idx_page + if layout := _get_csa_linear_tensor_layout(kv_cache_groups): + return layout.bytes_per_block if all( isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in kv_cache_groups ): @@ -1069,6 +1072,85 @@ def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bo return not kv_cache_spec +def _select_hybrid_kv_cache_group_size(layer_counts: Sequence[int]) -> int: + """Select a group size without letting a singleton dominate the layout.""" + min_num_layers = min(layer_counts) + max_num_layers = max(layer_counts) + if max_num_layers < min_num_layers * 1.5: + # If the number of layers is not much larger than the minimum number of + # layers, use the maximum number of layers as the group size to avoid + # too many padding layers. A typical example is gpt-oss-20b + eagle, + # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of + # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding + # layers while accommodating speculative decoding drafters that add + # extra layers to one attention type. + return max_num_layers + + # A singleton auxiliary state cache can otherwise force group_size=1 and + # multiply metadata preparation for every repeated cache type. Find the + # largest divisor shared by the repeated types that keeps total padding at + # or below 20%. Restrict this optimization to singleton outliers so existing + # hybrid layouts retain their current behavior. + if min_num_layers == 1: + repeated_counts = [count for count in layer_counts if count > 1] + if repeated_counts: + common_divisor = math.gcd(*repeated_counts) + total_layers = sum(layer_counts) + for candidate in range(common_divisor, 1, -1): + if common_divisor % candidate != 0: + continue + padded_layers = candidate * sum( + cdiv(count, candidate) for count in layer_counts + ) + if (padded_layers - total_layers) * 5 <= total_layers: + return candidate + + return min_num_layers + + +def _prefer_padding_sliding_window_layers( + same_type_layers: dict[KVCacheSpec, list[str]], group_size: int +) -> int: + """Prefer cheap sliding-window padding over full-context padding. + + A DFlash2 target can have 16 full-attention layers, 48 Mamba layers, and + only 5 draft sliding-window layers. Using the smallest bucket as the group + size pads the full-context buckets. Sliding-window state is bounded by its + window, so choose a common divisor of the other buckets when that moves no + more than one bucket's worth of padding onto each sliding-window bucket. + """ + sliding_window_counts = [ + len(layers) + for spec, layers in same_type_layers.items() + if isinstance(spec, (SlidingWindowSpec, SlidingWindowMLASpec)) + ] + other_counts = [ + len(layers) + for spec, layers in same_type_layers.items() + if not isinstance(spec, (SlidingWindowSpec, SlidingWindowMLASpec)) + ] + if ( + group_size <= 1 + or not sliding_window_counts + or not other_counts + or min(sliding_window_counts) != group_size + ): + return group_size + + common_divisor = math.gcd(*other_counts) + for candidate in range(common_divisor, group_size, -1): + if common_divisor % candidate != 0: + continue + if all((-count) % candidate <= count for count in sliding_window_counts): + logger.info( + "Using KV cache group size %d so bounded sliding-window " + "layers absorb padding instead of full-context layers", + candidate, + ) + return candidate + return group_size + + def _get_kv_cache_groups_uniform_page_size( kv_cache_spec: dict[str, KVCacheSpec], ) -> list[KVCacheGroupSpec]: @@ -1153,18 +1235,15 @@ def _get_kv_cache_groups_uniform_page_size( # is the minimum number of layers among all attention types. Need a better # strategy if we want to support more complex patterns (e.g., 20 full + 30 # sw, where the group size should be 10). - min_num_layers = min([len(layers) for layers in same_type_layers.values()]) - group_size = min_num_layers - max_num_layers = max([len(layers) for layers in same_type_layers.values()]) - if max_num_layers < min_num_layers * 1.5: - # If the number of layers is not much larger than the minimum number of - # layers, use the maximum number of layers as the group size to avoid - # too many padding layers. A typical example is gpt-oss-20b + eagle, - # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of - # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding - # layers while accommodating speculative decoding drafters that add - # extra layers to one attention type. - group_size = max_num_layers + layer_counts = [len(layers) for layers in same_type_layers.values()] + group_size = _select_hybrid_kv_cache_group_size(layer_counts) + group_size = _prefer_padding_sliding_window_layers(same_type_layers, group_size) + if group_size != min(layer_counts): + logger.info( + "Using hybrid KV cache group size %d for layer counts %s", + group_size, + sorted(layer_counts), + ) grouped_layers = [] for layers in same_type_layers.values(): num_padding_layers = group_size - len(layers) % group_size @@ -1248,6 +1327,244 @@ def _get_kv_cache_config_deepseek_v4( return num_blocks, kv_cache_tensors +def _pp_balanced_mamba_group_count( + vllm_config: VllmConfig, + mamba_layer_names: list[str], + mla_layer_names: list[str], +) -> int | None: + """Mamba group count for `_get_kv_cache_groups_glm5_next`. + + Under PP, each stage's largest projected Mamba group must fit its local MLA + count. Returns the smallest group count that satisfies every stage, or + ``None`` when a stage has Mamba layers but no MLA layer. + """ + num_groups = cdiv(len(mamba_layer_names), len(mla_layer_names)) + pp_size = vllm_config.parallel_config.pipeline_parallel_size + if pp_size == 1: + return num_groups + + from vllm.distributed.utils import get_pp_indices + from vllm.model_executor.models.utils import extract_layer_index + + total_layers = vllm_config.model_config.get_total_num_hidden_layers() + mamba_indices = [extract_layer_index(name) for name in mamba_layer_names] + mla_indices = [extract_layer_index(name) for name in mla_layer_names] + for rank in range(pp_size): + start, end = get_pp_indices(total_layers, rank, pp_size) + num_mamba = sum(start <= i < end for i in mamba_indices) + num_mla = sum(start <= i < end for i in mla_indices) + if not num_mamba: + continue + if not num_mla: + return None + # A stage's mamba layers are contiguous in round-robin order, so its + # largest slice under `count` groups is exactly cdiv(num_mamba, count), + # which fits num_mla slots iff count >= cdiv(num_mamba, num_mla). + num_groups = max(num_groups, cdiv(num_mamba, num_mla)) + return num_groups + + +def _get_kv_cache_groups_glm5_next( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Groups for GLM-5.3-Flash hybrids: MLA(+kpool indexer) + mamba. + + Mamba layers share the MLA layers' KV cache tensors. Indexer layers get + their own KV cache tensors as indexer blocks are small. + """ + mamba_specs = {k: v for k, v in kv_cache_spec.items() if isinstance(v, MambaSpec)} + tail_specs = { + k: v for k, v in kv_cache_spec.items() if isinstance(v, KpoolTailSpec) + } + attn_specs = { + k: v + for k, v in kv_cache_spec.items() + if not isinstance(v, (MambaSpec, KpoolTailSpec)) + } + if not mamba_specs or not all( + type(s) is MLAAttentionSpec for s in attn_specs.values() + ): + return None + mla_specs = cast(dict[str, MLAAttentionSpec], attn_specs) + idx_pages = {s.page_size_bytes for s in mla_specs.values() if s.compress_ratio > 1} + if not idx_pages: + return None + + assert all(s.page_size_padded is None for s in mla_specs.values()) + assert len(idx_pages) == 1 + mla_names = [n for n, s in mla_specs.items() if s.compress_ratio == 1] + mla_pages = {mla_specs[n].page_size_bytes for n in mla_names} + assert len(mla_pages) == 1 + mla_page = mla_pages.pop() + uniform_spec = UniformTypeKVCacheSpecs.from_specs(attn_specs) + assert uniform_spec is not None + + # Keep all indexer tails in one group and pad their pages to the indexer + # page size so each tail can share its sibling indexer's storage. + tail_group = None + if tail_specs: + idx_page = next(iter(idx_pages)) + padded_tail_specs: dict[str, KVCacheSpec] = { + name: replace(s, page_size_padded=idx_page) + for name, s in tail_specs.items() + } + tail_uniform = UniformTypeKVCacheSpecs.from_specs(padded_tail_specs) + assert tail_uniform is not None + tail_group = KVCacheGroupSpec(list(padded_tail_specs), tail_uniform) + + any_mamba = next(iter(mamba_specs.values())) + assert all(spec == any_mamba for spec in mamba_specs.values()) + if any_mamba.real_page_size_bytes > mla_page: + raise ValueError( + f"the mamba state page ({any_mamba.real_page_size_bytes} bytes) " + f"does not fit the MLA page ({mla_page} bytes); increase tensor " + "parallelism or use a wider KV cache dtype" + ) + padded_specs: dict[str, KVCacheSpec] = { + name: replace(any_mamba, page_size_padded=mla_page) for name in mamba_specs + } + num_groups = _pp_balanced_mamba_group_count( + vllm_config, list(mamba_specs), mla_names + ) + if num_groups is None: + raise ValueError( + "a pipeline stage has mamba layers but no MLA layer to share " + "slots with; realign the stage boundaries (VLLM_PP_LAYER_PARTITION)" + ) + mamba_grouped_names: list[list[str]] = [[] for _ in range(num_groups)] + for k, name in enumerate(mamba_specs): + mamba_grouped_names[k % num_groups].append(name) + return ( + [KVCacheGroupSpec(list(attn_specs), uniform_spec)] + + ([tail_group] if tail_group is not None else []) + + create_kv_cache_group_specs(padded_specs, mamba_grouped_names) + ) + + +def _glm5_next_tensor_layout( + kv_cache_groups: list[KVCacheGroupSpec], +) -> ( + tuple[ + KVCacheGroupSpec, + list[KVCacheGroupSpec], + list[str], + list[str], + int, + int, + list[str], + int, + ] + | None +): + """Detect the `_get_kv_cache_groups_glm5_next` layout from the + (possibly PP-projected) groups, so tensor emission and the accounting + paths can never disagree. + + Returns: + - (attn_group, mamba_groups, mla_names, idx_names, mla_page, idx_page, + tail_names, tail_page) + - None if the group config is invalid + """ + uniform_groups = [ + g + for g in kv_cache_groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + ] + mamba_groups = [ + g for g in kv_cache_groups if isinstance(g.kv_cache_spec, MambaSpec) + ] + # Both the MLA(+indexer) attn group and the kpool tail group are + # UniformTypeKVCacheSpecs; distinguish by the inner spec type. + attn_group: KVCacheGroupSpec | None = None + tail_group: KVCacheGroupSpec | None = None + for g in uniform_groups: + group_inner = cast(UniformTypeKVCacheSpecs, g.kv_cache_spec).kv_cache_specs + if all(type(s) is MLAAttentionSpec for s in group_inner.values()): + attn_group = g + elif all(isinstance(s, KpoolTailSpec) for s in group_inner.values()): + tail_group = g + if attn_group is None or not mamba_groups: + return None + if len(uniform_groups) + len(mamba_groups) != len(kv_cache_groups): + return None + attn_uniform = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec) + if not all( + type(s) is MLAAttentionSpec and s.page_size_padded is None + for s in attn_uniform.kv_cache_specs.values() + ): + return None + inner = cast(dict[str, MLAAttentionSpec], attn_uniform.kv_cache_specs) + mla_names = [n for n in attn_group.layer_names if inner[n].compress_ratio == 1] + idx_names = [n for n in attn_group.layer_names if inner[n].compress_ratio > 1] + mla_pages = {inner[n].page_size_bytes for n in mla_names} + idx_pages = {inner[n].page_size_bytes for n in idx_names} + if len(mla_pages) != 1 or len(idx_pages) != 1: + return None + mla_page = mla_pages.pop() + if any(g.kv_cache_spec.page_size_bytes != mla_page for g in mamba_groups): + return None + tail_names: list[str] = [] + tail_page = 0 + if tail_group is not None: + tail_names = list(tail_group.layer_names) + # The tail spec is padded to idx_page (co-owns the indexer tensor), so + # page_size_bytes returns idx_page. Callers need the *logical* tail page + # (2048 B) for transfer sizing and accounting; use unpadded. + tail_pages = { + cast(KpoolTailSpec, s).real_page_size_bytes + for s in cast( + UniformTypeKVCacheSpecs, tail_group.kv_cache_spec + ).kv_cache_specs.values() + } + if len(tail_pages) != 1: + return None + tail_page = tail_pages.pop() + return ( + attn_group, + mamba_groups, + mla_names, + idx_names, + mla_page, + idx_pages.pop(), + tail_names, + tail_page, + ) + + +def _get_kv_cache_config_csa_linear( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> tuple[int, list[KVCacheTensor]] | None: + layout = _get_csa_linear_tensor_layout(kv_cache_groups) + if layout is None: + return None + + num_blocks = available_memory // layout.bytes_per_block + num_blocks = may_override_num_blocks(vllm_config, num_blocks) + kv_cache_tensors = [ + KVCacheTensor( + size=layout.main_kv_page_size * num_blocks, + shared_by=[main_kv_name] + + [ + group.layer_names[index] + for group in layout.mamba_groups + if index < len(group.layer_names) + ], + ) + for index, main_kv_name in enumerate(layout.main_kv_names) + ] + kv_cache_tensors.extend( + KVCacheTensor( + size=layout.compressed_page_size * num_blocks, + shared_by=[compressed_name, layout.compressor_state_names[index]], + ) + for index, compressed_name in enumerate(layout.compressed_names) + ) + return num_blocks, kv_cache_tensors + + def get_kv_cache_config_from_groups( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], @@ -1292,6 +1609,43 @@ def get_kv_cache_config_from_groups( ) for layer_name in kv_cache_groups[0].layer_names ] + elif (glm5n := _glm5_next_tensor_layout(kv_cache_groups)) is not None: + ( + _, + mamba_groups, + mla_names, + idx_names, + mla_page, + idx_page, + tail_names, + _tail_page, + ) = glm5n + if tail_names: + assert len(idx_names) == len(tail_names) + per_block = len(mla_names) * mla_page + len(idx_names) * idx_page + num_blocks = may_override_num_blocks(vllm_config, available_memory // per_block) + kv_cache_tensors = [ + KVCacheTensor( + size=mla_page * num_blocks, + shared_by=[mla_name] + + [ + group.layer_names[i] + for group in mamba_groups + if i < len(group.layer_names) + ], + ) + for i, mla_name in enumerate(mla_names) + ] + [ + KVCacheTensor( + size=idx_page * num_blocks, + shared_by=([idx_name, tail_names[i]] if tail_names else [idx_name]), + ) + for i, idx_name in enumerate(idx_names) + ] + elif csa_config := _get_kv_cache_config_csa_linear( + vllm_config, kv_cache_groups, available_memory + ): + num_blocks, kv_cache_tensors = csa_config elif all( isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in kv_cache_groups @@ -1470,6 +1824,303 @@ def group_and_unify_kv_cache_specs( return [mla_uniform_spec, *swa_uniform_specs] +class _CSALinearCacheTuple(NamedTuple): + """Three cache owners attached to one compressed-sparse layer.""" + + layer_index: int + main_kv: tuple[str, FullAttentionSpec] + compressed: tuple[str, MLAAttentionSpec] + compressor_state: tuple[str, CircularBufferSpec] + + +@dataclass(frozen=True) +class _CSALinearTensorLayout: + main_kv_names: list[str] + compressed_names: list[str] + compressor_state_names: list[str] + mamba_groups: list[KVCacheGroupSpec] + main_kv_page_size: int + compressed_page_size: int + + @property + def bytes_per_block(self) -> int: + return len(self.main_kv_names) * ( + self.main_kv_page_size + self.compressed_page_size + ) + + +class _CSALinearRoles(NamedTuple): + main_kv: dict[str, FullAttentionSpec] + compressed: dict[str, MLAAttentionSpec] + compressor_state: dict[str, CircularBufferSpec] + mamba: dict[str, MambaSpec] + + +def _classify_csa_linear_specs( + kv_cache_spec: Mapping[str, KVCacheSpec], +) -> _CSALinearRoles | None: + if not any(type(spec) is CircularBufferSpec for spec in kv_cache_spec.values()): + return None + + roles = _CSALinearRoles(main_kv={}, compressed={}, compressor_state={}, mamba={}) + unsupported: list[str] = [] + for name, spec in kv_cache_spec.items(): + if type(spec) is FullAttentionSpec: + roles.main_kv[name] = spec + elif type(spec) is MLAAttentionSpec and spec.compress_ratio > 1: + roles.compressed[name] = spec + elif type(spec) is CircularBufferSpec: + roles.compressor_state[name] = spec + elif type(spec) is MambaSpec: + roles.mamba[name] = spec + else: + unsupported.append(name) + if unsupported: + raise ValueError(f"CSA+linear unsupported cache owners: {unsupported}.") + if not roles.main_kv or not roles.compressed or not roles.mamba: + raise ValueError( + "CSA+linear requires main KV, compressed, compressor-state, and " + "Mamba cache owners." + ) + return roles + + +def _get_csa_linear_cache_tuples( + roles: _CSALinearRoles, +) -> list[_CSALinearCacheTuple]: + from vllm.model_executor.models.utils import extract_layer_index + + def by_layer_index( + specs: Mapping[str, KVCacheSpec], role: str + ) -> dict[int, tuple[str, KVCacheSpec]]: + indexed: dict[int, tuple[str, KVCacheSpec]] = {} + for name, spec in specs.items(): + try: + layer_index = extract_layer_index(name) + except (AssertionError, ValueError) as exc: + raise ValueError( + f"CSA+linear {role} owner {name!r} does not identify " + "exactly one transformer layer." + ) from exc + if layer_index in indexed: + other_name = indexed[layer_index][0] + raise ValueError( + f"CSA+linear layer {layer_index} has duplicate {role} " + f"owners: {other_name!r} and {name!r}." + ) + indexed[layer_index] = (name, spec) + return indexed + + main_kv = by_layer_index(roles.main_kv, "main-KV") + compressed = by_layer_index(roles.compressed, "compressed") + compressor_state = by_layer_index(roles.compressor_state, "compressor-state") + layer_indices = set(main_kv) + if set(compressed) != layer_indices or set(compressor_state) != layer_indices: + raise ValueError( + "CSA+linear main-KV, compressed, and compressor-state owners must " + "have matching transformer-layer indices." + ) + + return [ + _CSALinearCacheTuple( + layer_index, + cast(tuple[str, FullAttentionSpec], main_kv[layer_index]), + cast(tuple[str, MLAAttentionSpec], compressed[layer_index]), + cast(tuple[str, CircularBufferSpec], compressor_state[layer_index]), + ) + for layer_index in sorted(layer_indices) + ] + + +def _get_csa_linear_mamba_group_count( + vllm_config: VllmConfig, + mamba_names: list[str], + main_kv_names: list[str], +) -> int | None: + num_groups = cdiv(len(mamba_names), len(main_kv_names)) + pp_size = vllm_config.parallel_config.pipeline_parallel_size + if pp_size == 1: + return num_groups + + from vllm.distributed.utils import get_pp_indices + from vllm.model_executor.models.utils import extract_layer_index + + total_layers = vllm_config.model_config.get_total_num_hidden_layers() + mamba_indices = [extract_layer_index(name) for name in mamba_names] + main_kv_indices = [extract_layer_index(name) for name in main_kv_names] + for rank in range(pp_size): + start, end = get_pp_indices(total_layers, rank, pp_size) + num_mamba = sum(start <= index < end for index in mamba_indices) + num_main_kv = sum(start <= index < end for index in main_kv_indices) + if not num_mamba: + continue + if not num_main_kv: + return None + num_groups = max(num_groups, cdiv(num_mamba, num_main_kv)) + return num_groups + + +def _get_kv_cache_groups_csa_linear( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Build Qwen4Exp's main/compressed/ring/recurrent cache groups.""" + + roles = _classify_csa_linear_specs(kv_cache_spec) + if roles is None: + return None + tuples = _get_csa_linear_cache_tuples(roles) + + expected_local_kv_heads = vllm_config.model_config.get_num_kv_heads( + vllm_config.parallel_config + ) + if any( + cache.main_kv[1].num_kv_heads != expected_local_kv_heads for cache in tuples + ): + raise ValueError( + "CSA+linear main-KV specs do not match the TP-local KV-head geometry." + ) + + for cache in tuples: + _, main_kv = cache.main_kv + _, compressed = cache.compressed + _, compressor_state = cache.compressor_state + if not ( + main_kv.block_size == compressed.block_size + and compressor_state.real_page_size_bytes <= compressed.page_size_bytes + and all( + spec.page_size_padded is None + for spec in (main_kv, compressed, compressor_state) + ) + ): + raise ValueError( + f"CSA+linear layer {cache.layer_index} violates cache geometry." + ) + + shapes = { + ( + cache.compressed[1].compress_ratio, + cache.main_kv[1].block_size, + cache.main_kv[1].page_size_bytes, + cache.compressed[1].page_size_bytes, + ) + for cache in tuples + } + if len(shapes) != 1: + raise ValueError( + "CSA+linear layers must share one block size, compression ratio, " + "and main/compressed page geometry." + ) + _, _, main_kv_page, compressed_page = next(iter(shapes)) + + padded_compressor_specs: dict[str, KVCacheSpec] = { + cache.compressor_state[0]: replace( + cache.compressor_state[1], page_size_padded=compressed_page + ) + for cache in tuples + } + compressed_sparse_specs: dict[str, KVCacheSpec] = { + name: spec + for cache in tuples + for name, spec in (cache.main_kv, cache.compressed) + } + compressed_sparse_uniform = UniformTypeKVCacheSpecs.from_specs( + compressed_sparse_specs + ) + compressor_uniform = UniformTypeKVCacheSpecs.from_specs(padded_compressor_specs) + if compressed_sparse_uniform is None or compressor_uniform is None: + raise ValueError("CSA+linear cache owners have incompatible lifetimes.") + + groups = [ + KVCacheGroupSpec(list(compressed_sparse_specs), compressed_sparse_uniform), + KVCacheGroupSpec(list(padded_compressor_specs), compressor_uniform), + ] + main_kv_names = [cache.main_kv[0] for cache in tuples] + for tp_replicated in (False, True): + names = [ + name + for name, spec in roles.mamba.items() + if spec.tp_replicated is tp_replicated + ] + if not names: + continue + representative = roles.mamba[names[0]] + if any(roles.mamba[name] != representative for name in names[1:]): + policy = "replicated" if tp_replicated else "sharded" + raise ValueError( + f"CSA+linear {policy} Mamba owners must use one cache spec." + ) + unpadded_page = replace(representative, page_size_padded=None).page_size_bytes + if unpadded_page > main_kv_page: + raise ValueError( + f"CSA+linear Mamba owner {names[0]!r} needs {unpadded_page} " + f"bytes, but a main-KV page has {main_kv_page} bytes." + ) + num_groups = _get_csa_linear_mamba_group_count( + vllm_config, names, main_kv_names + ) + if num_groups is None: + raise ValueError( + "CSA+linear pipeline stage has Mamba owners but no main-KV slots." + ) + padded_spec = replace(representative, page_size_padded=main_kv_page) + grouped_names: list[list[str]] = [[] for _ in range(num_groups)] + for index, name in enumerate(names): + grouped_names[index % num_groups].append(name) + groups.extend( + KVCacheGroupSpec(group_names, padded_spec) for group_names in grouped_names + ) + return groups + + +def _get_csa_linear_tensor_layout( + kv_cache_groups: list[KVCacheGroupSpec], +) -> _CSALinearTensorLayout | None: + compressed_sparse: Mapping[str, KVCacheSpec] | None = None + compressor_state: Mapping[str, KVCacheSpec] | None = None + mamba_groups: list[KVCacheGroupSpec] = [] + for group in kv_cache_groups: + if not group.layer_names: + continue + spec = group.kv_cache_spec + if isinstance(spec, MambaSpec): + mamba_groups.append(group) + continue + if not isinstance(spec, UniformTypeKVCacheSpecs): + return None + member = next(iter(spec.kv_cache_specs.values())) + if type(member) is CircularBufferSpec: + compressor_state = spec.kv_cache_specs + elif type(member) is FullAttentionSpec: + compressed_sparse = spec.kv_cache_specs + else: + return None + if compressed_sparse is None or compressor_state is None: + return None + + main_kv_names = [ + name + for name, spec in compressed_sparse.items() + if type(spec) is FullAttentionSpec + ] + compressed_names = [ + name + for name, spec in compressed_sparse.items() + if type(spec) is MLAAttentionSpec + ] + if not main_kv_names or not compressed_names: + return None + + return _CSALinearTensorLayout( + main_kv_names=main_kv_names, + compressed_names=compressed_names, + compressor_state_names=list(compressor_state), + mamba_groups=mamba_groups, + main_kv_page_size=compressed_sparse[main_kv_names[0]].page_size_bytes, + compressed_page_size=compressed_sparse[compressed_names[0]].page_size_bytes, + ) + + def _select_deepseek_v4_tuple_width( vllm_config: VllmConfig, grouped_specs: list[UniformTypeKVCacheSpecs], @@ -1663,6 +2314,12 @@ def get_kv_cache_groups( ) _annotate_eagle_groups_deepseek_v4(vllm_config, kv_cache_spec, kv_cache_groups) return kv_cache_groups + elif glm5_groups := _get_kv_cache_groups_glm5_next(vllm_config, kv_cache_spec): + # GLM-5.3-Flash case: one uniform MLA(+kpool indexer) attention group + # plus mamba groups whose layers share the MLA layers' slot tensors. + return glm5_groups + elif csa_groups := _get_kv_cache_groups_csa_linear(vllm_config, kv_cache_spec): + return csa_groups # Pull HiddenStateCacheSpec layers out before the general multi-group # path so they don't affect page-size unification or grouping. @@ -1770,6 +2427,16 @@ def _max_memory_usage_bytes_from_groups( spec.max_memory_usage_bytes(vllm_config) for spec in per_layer_specs.values() ) + elif layout := _get_csa_linear_tensor_layout(kv_cache_groups): + blocks_needed = sum( + cdiv( + group.kv_cache_spec.max_memory_usage_bytes(vllm_config), + group.kv_cache_spec.page_size_bytes, + ) + for group in kv_cache_groups + if group.layer_names + ) + return layout.bytes_per_block * blocks_needed elif all( isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in kv_cache_groups @@ -1796,6 +2463,36 @@ def _max_memory_usage_bytes_from_groups( total_max_mem_usage_bytes += g_max_mem_usage_page_bytes return total_max_mem_usage_bytes + elif (glm5n := _glm5_next_tensor_layout(kv_cache_groups)) is not None: + # GLM-5.3-Flash hybrid slot sharing: every block id — attention-, mamba-, + # or tail-owned — is charged the full per-block byte sum (mla + idx). + # The tail co-owns the indexer tensor (1 block/req via KpoolTailManager), + # so it adds to the shared block-id demand like mamba, not as a separate + # tensor. + ( + attn_group, + mamba_groups, + mla_names, + idx_names, + mla_page, + idx_page, + tail_names, + _tail_page, + ) = glm5n + uniform_spec = attn_group.kv_cache_spec + assert isinstance(uniform_spec, UniformTypeKVCacheSpecs) + blocks_needed = uniform_spec.max_memory_usage_pages(vllm_config) + for group in mamba_groups: + spec = group.kv_cache_spec + blocks_needed += cdiv( + spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes + ) + if tail_names: + # Tail: 1 block/req (KpoolTailSpec.max_admission_blocks_per_request + # == 1), drawn from the shared pool. + blocks_needed += 1 + return blocks_needed * (len(mla_names) * mla_page + len(idx_names) * idx_page) + # General case: group_size pools, each shared by one layer per group # Memory = group_size * page_size * blocks_for_max_len group_size = max(len(group.layer_names) for group in kv_cache_groups) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 86b5a1aeff..fb27764898 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1646,17 +1646,25 @@ def update_from_output( request.status = RequestStatus.FINISHED_STOPPED stopped = True - if new_token_ids and self.structured_output_manager.should_advance(request): + if new_token_ids and self.structured_output_manager.should_advance( + request, new_token_ids=new_token_ids + ): struct_output_request = request.structured_output_request assert struct_output_request is not None - assert struct_output_request.grammar is not None - if not struct_output_request.grammar.accept_tokens( # type: ignore[union-attr] - req_id, new_token_ids + grammar = struct_output_request.grammar + assert grammar is not None + advance_token_ids = ( + self.structured_output_manager.trim_reasoning_for_advance( + request, new_token_ids + ) + ) + if advance_token_ids and not grammar.accept_tokens( + req_id, advance_token_ids ): logger.error( "Unexpected: grammar rejected tokens %s for request %s. " "Terminating request.", - new_token_ids, + advance_token_ids, req_id, ) request.status = RequestStatus.FINISHED_ERROR diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index b3980c592f..2a76eb9963 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -14,9 +14,11 @@ ) from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, + CircularBufferSpec, CrossAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, + KpoolTailSpec, KVCacheSpec, MambaSpec, MLAAttentionSpec, @@ -592,6 +594,98 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int: return num_common_blocks +class CircularBufferManager(FullAttentionManager): + """Own exactly one QSA compressor-ring block per live request.""" + + def _claim_ring_block(self, request_id: str) -> list[KVCacheBlock]: + req_blocks = self.req_to_blocks[request_id] + if req_blocks: + return [] + new_blocks = self.block_pool.get_new_blocks(1) + req_blocks.extend(new_blocks) + self.new_block_ids.extend(block.block_id for block in new_blocks) + return new_blocks + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + del ( + num_tokens, + new_computed_blocks, + total_computed_tokens, + num_tokens_main_model, + apply_admission_cap, + ) + return 0 if self.req_to_blocks.get(request_id) else 1 + + def allocate_new_computed_blocks( + self, + request_id: str, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + del new_computed_blocks, num_local_computed_tokens, num_external_computed_tokens + self._claim_ring_block(request_id) + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + del num_tokens, num_tokens_main_model + return self._claim_ring_block(request_id) + + @classmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + use_eagle: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[list[KVCacheBlock], ...]: + del ( + block_hashes, + max_length, + block_pool, + kv_cache_spec, + use_eagle, + alignment_tokens, + dcp_world_size, + pcp_world_size, + ) + return tuple([] for _ in kv_cache_group_ids) + + def cache_blocks( + self, + request: Request, + num_tokens: int, + alignment_tokens: int | None = None, + ) -> None: + del request, num_tokens, alignment_tokens + + def remove_skipped_blocks( + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: + del request_id, processed_computed_tokens, num_prompt_tokens + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + del running_request_id + return 0 + + class PrefixAnchoredSWAManager(FullAttentionManager): """KV cache manager for prefix-anchored sliding-window attention. @@ -789,6 +883,104 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int: return 0 +class KpoolTailManager(FullAttentionManager): + """One non-shareable circular tail block per request.""" + + @classmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + use_eagle: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[list[KVCacheBlock], ...]: + del ( + block_hashes, + max_length, + block_pool, + kv_cache_spec, + use_eagle, + alignment_tokens, + dcp_world_size, + pcp_world_size, + ) + return tuple([] for _ in kv_cache_group_ids) + + def cache_blocks( + self, + request: Request, + num_tokens: int, + alignment_tokens: int | None = None, + ) -> None: + del request, num_tokens, alignment_tokens + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + del running_request_id + return 0 + + def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: + del num_computed_tokens + return 0 + + def remove_skipped_blocks( + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: + del request_id, processed_computed_tokens, num_prompt_tokens + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + del ( + num_tokens, + new_computed_blocks, + total_computed_tokens, + num_tokens_main_model, + apply_admission_cap, + ) + return max(1 - len(self.req_to_blocks.get(request_id, ())), 0) + + def _allocate_tail_block(self, request_id: str) -> list[KVCacheBlock]: + req_blocks = self.req_to_blocks[request_id] + if req_blocks: + return [] + new_blocks = self.block_pool.get_new_blocks(1) + req_blocks.extend(new_blocks) + self.new_block_ids.extend(block.block_id for block in new_blocks) + return new_blocks + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + del num_tokens, num_tokens_main_model + return self._allocate_tail_block(request_id) + + def allocate_new_computed_blocks( + self, + request_id: str, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + if new_computed_blocks: + raise ValueError("KPool tail blocks cannot be prefix-cache hits.") + if num_local_computed_tokens or num_external_computed_tokens: + self._allocate_tail_block(request_id) + + class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec, **kwargs) -> None: super().__init__(kv_cache_spec, **kwargs) @@ -1322,10 +1514,12 @@ def __init__( PrefixAnchoredSWASpec: PrefixAnchoredSWAManager, SlidingWindowSpec: SlidingWindowManager, SlidingWindowMLASpec: SlidingWindowManager, + KpoolTailSpec: KpoolTailManager, ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager, MambaSpec: MambaManager, CrossAttentionSpec: CrossAttentionManager, SinkFullAttentionSpec: SinkFullAttentionManager, + CircularBufferSpec: CircularBufferManager, } diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 4e094eb6a6..9210b37eec 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -282,8 +282,14 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups if kv_cache_groups: + participating_block_sizes = [ + group.kv_cache_spec.block_size + for group in kv_cache_groups + if group.kv_cache_spec.prefix_cacheable + ] vllm_config.cache_config.block_size = min( - g.kv_cache_spec.block_size for g in kv_cache_groups + participating_block_sizes + or [group.kv_cache_spec.block_size for group in kv_cache_groups] ) vllm_config.validate_block_size() @@ -461,9 +467,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: profile_sample_ms = 0.0 profile_update_ms = 0.0 if profile_ddtree_engine: - profile_step = getattr( - self, "_dflash_ddtree_engine_profile_step", 0 - ) + 1 + profile_step = getattr(self, "_dflash_ddtree_engine_profile_step", 0) + 1 self._dflash_ddtree_engine_profile_step = profile_step profile_t0 = time.perf_counter() @@ -478,9 +482,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: profile_part_t0 = time.perf_counter() future = self.model_executor.execute_model(scheduler_output, non_block=True) if profile_ddtree_engine: - profile_execute_submit_ms = ( - time.perf_counter() - profile_part_t0 - ) * 1000.0 + profile_execute_submit_ms = (time.perf_counter() - profile_part_t0) * 1000.0 profile_part_t0 = time.perf_counter() grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output) if profile_ddtree_engine: @@ -497,9 +499,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: profile_part_t0 = time.perf_counter() if profile_ddtree_engine else 0.0 model_output = self.model_executor.sample_tokens(grammar_output) if profile_ddtree_engine: - profile_sample_ms = ( - time.perf_counter() - profile_part_t0 - ) * 1000.0 + profile_sample_ms = (time.perf_counter() - profile_part_t0) * 1000.0 # Before processing the model output, process any aborts that happened # during the model execution. @@ -611,9 +611,7 @@ def step_with_batch_queue( trace_schedule_t0 = time.perf_counter() if trace_log else 0.0 scheduler_output = self.scheduler.schedule() if trace_log: - trace_schedule_ms = ( - time.perf_counter() - trace_schedule_t0 - ) * 1000.0 + trace_schedule_ms = (time.perf_counter() - trace_schedule_t0) * 1000.0 trace_scheduled_tokens = scheduler_output.total_num_scheduled_tokens with self.log_error_detail(scheduler_output): trace_execute_t0 = time.perf_counter() if trace_log else 0.0 diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 83f32dfccf..c9999df01e 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -100,6 +100,11 @@ class KVCacheSpec: # number of tokens in a block block_size: int + @property + def prefix_cacheable(self) -> bool: + """Whether blocks owned by this spec can be shared by a prefix.""" + return True + @property def page_size_bytes(self) -> int: """ @@ -123,6 +128,10 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: """ raise NotImplementedError + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: + del vllm_config + return cdiv(max_len, self.block_size) + def copy_with_new_block_size(self, block_size: int) -> Self: """ Create a new KVCacheSpec from self but replacing the block size. @@ -352,6 +361,12 @@ def storage_block_size(self) -> int: @property def real_page_size_bytes(self) -> int: + if self.model_version == "glm5_next" and self.cache_dtype_str in { + "fp8", + "fp8_e4m3", + }: + # GLM-5.3 NoPE MLA: 512 E4M3 bytes + eight UE8M0 group scales. + return self.storage_block_size * 520 if self.cache_dtype_str == "fp8_ds_mla": if self.model_version == "deepseek_v4": # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token. @@ -538,6 +553,30 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: return max_blocks * self.page_size_bytes +@dataclass(frozen=True, kw_only=True) +class CircularBufferSpec(AttentionSpec): + """One fixed block per request for QSA's uncompressed key ring.""" + + head_size_v: int = 0 + + @property + def real_page_size_bytes(self) -> int: + return ( + self.block_size + * self.num_kv_heads + * (self.head_size + self.head_size_v) + * get_dtype_size(self.dtype) + ) + + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: + del vllm_config + return self.page_size_bytes + + @property + def prefix_cacheable(self) -> bool: + return False + + @dataclass(frozen=True, kw_only=True) class SlidingWindowMLASpec(SlidingWindowSpec): """Sliding window attention with MLA cache format.""" @@ -603,21 +642,50 @@ def merge(cls, specs: list[Self]) -> Self: ) +@dataclass(frozen=True, kw_only=True) +class KpoolTailSpec(SlidingWindowSpec): + """Per-request circular cache for an incomplete GLM KPool group.""" + + def max_admission_blocks_per_request( + self, max_num_batched_tokens: int, max_model_len: int + ) -> int: + del max_num_batched_tokens, max_model_len + return 1 + + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: + return 1 + + def is_uniform_with_collection( + self, kv_cache_specs: dict[str, KVCacheSpec] + ) -> bool: + return all(isinstance(spec, KpoolTailSpec) for spec in kv_cache_specs.values()) + + @property + def prefix_cacheable(self) -> bool: + return False + + @dataclass(frozen=True) class MambaSpec(KVCacheSpec): shapes: tuple[tuple[int, ...], ...] - dtypes: tuple[torch.dtype] + dtypes: tuple[torch.dtype, ...] page_size_padded: int | None = None mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2 mamba_cache_mode: str = "none" num_speculative_blocks: int = 0 + # PLE short-conv state is replicated; GDN state is TP-sharded. + tp_replicated: bool = False @property - def page_size_bytes(self) -> int: - page_size = sum( + def real_page_size_bytes(self) -> int: + return sum( prod(shape) * get_dtype_size(dtype) for (shape, dtype) in zip(self.shapes, self.dtypes) ) + + @property + def page_size_bytes(self) -> int: + page_size = self.real_page_size_bytes if self.page_size_padded is not None: assert self.page_size_padded >= page_size return self.page_size_padded @@ -718,6 +786,10 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): kv_cache_specs: dict[str, KVCacheSpec] + @property + def prefix_cacheable(self) -> bool: + return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values()) + @property def page_size_bytes(self) -> int: return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values()) @@ -749,6 +821,10 @@ def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool: and spec.sliding_window == one_spec.sliding_window for spec in kv_cache_specs.values() ) + elif isinstance(one_spec, CircularBufferSpec): + return all( + isinstance(spec, CircularBufferSpec) for spec in kv_cache_specs.values() + ) elif isinstance(one_spec, FullAttentionSpec): return all( isinstance(spec, FullAttentionSpec) for spec in kv_cache_specs.values() diff --git a/vllm/v1/sample/thinking_budget_state.py b/vllm/v1/sample/thinking_budget_state.py index 4dd764766d..0f7ab58187 100644 --- a/vllm/v1/sample/thinking_budget_state.py +++ b/vllm/v1/sample/thinking_budget_state.py @@ -153,9 +153,7 @@ def update_state( # the suffix that is present; stripping the full spec length in # the repeated layout drops one real output token. spec_suffix_len = ( - max(0, spec_len - 1) - if last_row_for_req is not None - else spec_len + max(0, spec_len - 1) if last_row_for_req is not None else spec_len ) # Only strip draft suffix when there are spec tokens; ``[:-0]`` would # clear the whole list (Python treats stop index 0 as "up to empty"). @@ -163,9 +161,7 @@ def update_state( spec_suffix_len > 0 and len(state["output_tok_ids"]) >= spec_suffix_len ): - state["output_tok_ids"] = state["output_tok_ids"][ - :-spec_suffix_len - ] + state["output_tok_ids"] = state["output_tok_ids"][:-spec_suffix_len] self._update_think_state(state) def apply_to_logits( diff --git a/vllm/v1/spec_decode/ddtree_metadata.py b/vllm/v1/spec_decode/ddtree_metadata.py index 0cf0cf2d07..e8dc694676 100644 --- a/vllm/v1/spec_decode/ddtree_metadata.py +++ b/vllm/v1/spec_decode/ddtree_metadata.py @@ -52,10 +52,10 @@ def from_tree(cls, *, prompt_len: int, tree: DDTree) -> DDTreeVerifierMetadata: tree_token_ids = tree.token_ids_for_verifier() parent_indices = tree.parent_indices_for_verifier() node_depths = tree.node_depths_for_verifier() - tree_position_ids = tuple(prompt_len + depth - 1 - for depth in node_depths) + tree_position_ids = tuple(prompt_len + depth - 1 for depth in node_depths) compact_logits_indices = (prompt_len - 1,) + tuple( - prompt_len + offset for offset in range(len(tree_token_ids))) + prompt_len + offset for offset in range(len(tree_token_ids)) + ) edge_parent_compact_indices: list[int] = [] node_compact_indices: list[int] = [] @@ -98,8 +98,7 @@ def _offset_parent_indices( parent_indices: tuple[int, ...], tree_start: int, ) -> list[int]: - return [parent if parent < 0 else parent + tree_start - for parent in parent_indices] + return [parent if parent < 0 else parent + tree_start for parent in parent_indices] def make_prefill_tree_attention_mask( @@ -115,7 +114,7 @@ def make_prefill_tree_attention_mask( visible = torch.zeros((total_len, total_len), device=device, dtype=torch.bool) for row in range(prompt_len): - visible[row, :row + 1] = True + visible[row, : row + 1] = True for node in tree.non_root_nodes: row = prompt_len + node.index - 1 @@ -126,10 +125,9 @@ def make_prefill_tree_attention_mask( col = prompt_len + ancestor_index - 1 visible[row, col] = True - mask = torch.full((total_len, total_len), - torch.finfo(dtype).min, - device=device, - dtype=dtype) + mask = torch.full( + (total_len, total_len), torch.finfo(dtype).min, device=device, dtype=dtype + ) mask.masked_fill_(visible, 0) return mask.unsqueeze(0).unsqueeze(0) @@ -157,7 +155,8 @@ def make_batched_metadata( max_total_len = max( prompt_len + len(tree.non_root_nodes) - for prompt_len, tree in zip(prompt_lens, trees, strict=True)) + for prompt_len, tree in zip(prompt_lens, trees, strict=True) + ) tree_token_ids: list[int] = [] parent_indices: list[int] = [] node_depths: list[int] = [] @@ -172,15 +171,18 @@ def make_batched_metadata( for metadata in per_request: tree_token_ids.extend(metadata.tree_token_ids) parent_indices.extend( - _offset_parent_indices(metadata.parent_indices, tree_start)) + _offset_parent_indices(metadata.parent_indices, tree_start) + ) node_depths.extend(metadata.node_depths) compact_logits_indices.extend( - seq_start + index for index in metadata.compact_logits_indices) + seq_start + index for index in metadata.compact_logits_indices + ) edge_parent_compact_indices.extend( - compact_start + index - for index in metadata.edge_parent_compact_indices) + compact_start + index for index in metadata.edge_parent_compact_indices + ) node_compact_indices.extend( - compact_start + index for index in metadata.node_compact_indices) + compact_start + index for index in metadata.node_compact_indices + ) row = list(metadata.all_position_ids()) row.extend([0] * (max_total_len - len(row))) @@ -192,24 +194,20 @@ def make_batched_metadata( return BatchedDDTreeVerifierMetadata( prompt_lens=tuple(prompt_lens), - tree_token_ids=torch.tensor(tree_token_ids, device=device, - dtype=torch.int32), - parent_indices=torch.tensor(parent_indices, device=device, - dtype=torch.int32), + tree_token_ids=torch.tensor(tree_token_ids, device=device, dtype=torch.int32), + parent_indices=torch.tensor(parent_indices, device=device, dtype=torch.int32), node_depths=torch.tensor(node_depths, device=device, dtype=torch.int32), - position_ids=torch.tensor(position_rows, device=device, - dtype=torch.long), - cu_num_tree_nodes=torch.tensor(cu_nodes, device=device, - dtype=torch.int32), - compact_logits_indices=torch.tensor(compact_logits_indices, - device=device, - dtype=torch.int32), - edge_parent_compact_indices=torch.tensor(edge_parent_compact_indices, - device=device, - dtype=torch.int32), - node_compact_indices=torch.tensor(node_compact_indices, - device=device, - dtype=torch.int32), + position_ids=torch.tensor(position_rows, device=device, dtype=torch.long), + cu_num_tree_nodes=torch.tensor(cu_nodes, device=device, dtype=torch.int32), + compact_logits_indices=torch.tensor( + compact_logits_indices, device=device, dtype=torch.int32 + ), + edge_parent_compact_indices=torch.tensor( + edge_parent_compact_indices, device=device, dtype=torch.int32 + ), + node_compact_indices=torch.tensor( + node_compact_indices, device=device, dtype=torch.int32 + ), ) @@ -226,7 +224,8 @@ def greedy_sample_from_compact_logits( if compact_logits.shape[0] != expected_rows: raise ValueError( f"compact_logits row mismatch: expected {expected_rows}, " - f"got {compact_logits.shape[0]}") + f"got {compact_logits.shape[0]}" + ) target_argmax = compact_logits.argmax(dim=-1) path_to_compact_index: dict[tuple[int, ...], int] = {(): 0} @@ -252,7 +251,8 @@ def greedy_sample_from_compact_top_tokens( if compact_top_tokens.shape[0] != expected_rows: raise ValueError( f"compact_top_tokens row mismatch: expected {expected_rows}, " - f"got {compact_top_tokens.shape[0]}") + f"got {compact_top_tokens.shape[0]}" + ) target_tokens = compact_top_tokens.detach().cpu().tolist() path_to_compact_index: dict[tuple[int, ...], int] = {(): 0} diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 98bed3ac9d..a643262407 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -231,7 +231,7 @@ def __init__( # We need to get the hidden size from the draft model config because # the draft model's hidden size can be different from the target model's # hidden size (e.g., Llama 3.3 70B). - self.hidden_size = self.draft_model_config.get_hidden_size() + self.hidden_size = self._get_hidden_size() self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size() # DeepSeek V4 MTP consumes the target's pre-hc_head residual stream, @@ -461,6 +461,10 @@ def __init__( self.allowed_attn_types = tuple(rocm_types) + def _get_hidden_size(self) -> int: + """Return the hidden width consumed by the draft model.""" + return self.draft_model_config.get_hidden_size() + def _raise_if_padded_drafter_batch_disabled(self): if self.speculative_config.disable_padded_drafter_batch: raise NotImplementedError( diff --git a/vllm/v1/spec_decode/qwen4_exp.py b/vllm/v1/spec_decode/qwen4_exp.py new file mode 100644 index 0000000000..a1792b2c27 --- /dev/null +++ b/vllm/v1/spec_decode/qwen4_exp.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from copy import copy + +import torch + +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.worker.utils import AttentionGroup + + +class Qwen4ExpMTPProposer(EagleProposer): + """Speculative decoding proposer for Qwen4Exp MTP.""" + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + runner=None, + ) -> None: + super().__init__(vllm_config, device, runner) + self._per_group_block_tables: dict[int, torch.Tensor] = {} + + def _get_hidden_size(self) -> int: + """Return the multi-stream feedback width consumed by Qwen MTP.""" + return int( + self.draft_model_config.get_hidden_size() + * self.draft_model_config.hf_config.hc_mult + ) + + def model_returns_tuple(self) -> bool: + """Qwen MTP returns separate logits and feedback hidden states.""" + return True + + def set_per_group_block_table( + self, + gid: int, + block_table: torch.Tensor, + ) -> None: + """Stage one scheduler group's block table for drafting.""" + self._per_group_block_tables[gid] = block_table + + def build_per_group_and_layer_attn_metadata( + self, + common_attn_metadata: CommonAttentionMetadata, + draft_index: int = 0, + ) -> tuple[list[object], dict[str, object]]: + """Build each Qwen cache owner's metadata from its scheduler group.""" + per_group_attn_metadata: list[object] = [] + per_layer_attn_metadata: dict[str, object] = {} + common_by_gid: dict[int, CommonAttentionMetadata] = {} + num_reqs = common_attn_metadata.num_reqs + + for attn_group in self.draft_attn_groups: + gid = attn_group.kv_cache_group_id + group_common = common_by_gid.get(gid) + if group_common is None: + if gid == self.kv_cache_gid: + group_common = common_attn_metadata + else: + block_table = self._per_group_block_tables.get(gid) + assert block_table is not None, ( + f"Missing Qwen draft block table for KV cache group {gid}" + ) + group_common = copy(common_attn_metadata) + group_common.block_table_tensor = block_table[:num_reqs] + common_by_gid[gid] = group_common + + attn_metadata = attn_group.get_metadata_builder().build_for_drafting( + common_attn_metadata=group_common, + draft_index=draft_index, + ) + per_group_attn_metadata.append(attn_metadata) + for layer_name in attn_group.layer_names: + per_layer_attn_metadata[layer_name] = attn_metadata + + return per_group_attn_metadata, per_layer_attn_metadata + + def initialize_attn_backend( + self, + kv_cache_config: KVCacheConfig, + kernel_block_sizes: list[int] | None = None, + ) -> None: + """Initialize Qwen main, compressed, and circular cache owners.""" + num_mtp_layers = self.draft_model_config.hf_text_config.mtp_num_hidden_layers + if num_mtp_layers != 1: + raise NotImplementedError( + "Qwen4Exp MTP proposer only supports one MTP layer" + ) + assert kernel_block_sizes is not None, ( + "Qwen MTP requires resolved kernel block sizes" + ) + assert len(kernel_block_sizes) == len(kv_cache_config.kv_cache_groups), ( + "Qwen MTP requires one kernel block size per KV cache group" + ) + + all_attn_layers = get_layers_from_vllm_config( + self.vllm_config, + AttentionLayerBase, # type: ignore[type-abstract] + ) + layer_to_gid, layer_to_spec = self._map_draft_layers_to_groups(kv_cache_config) + main_layers = [ + name + for name, spec in layer_to_spec.items() + if type(spec) is FullAttentionSpec + ] + assert len(main_layers) == 1, ( + "Qwen4Exp MTP requires exactly one main cache owner" + ) + self.kv_cache_gid = layer_to_gid[main_layers[0]] + + attention_groups: list[AttentionGroup] = [] + for layer_name in sorted(self._draft_attn_layer_names): + attn_layer = all_attn_layers[layer_name] + gid = layer_to_gid[layer_name] + attn_group = AttentionGroup( + backend=attn_layer.get_attn_backend(), + layer_names=[layer_name], + kv_cache_spec=layer_to_spec[layer_name], + kv_cache_group_id=gid, + ) + attn_group.create_metadata_builders( + self.vllm_config, + self.device, + kernel_block_size=kernel_block_sizes[gid], + ) + attention_groups.append(attn_group) + + self.draft_attn_groups = sorted( + attention_groups, + key=lambda group: ( + group.kv_cache_group_id != self.kv_cache_gid, + group.kv_cache_group_id, + group.backend.full_cls_name(), + group.layer_names[0], + ), + ) + self.block_size = kernel_block_sizes[self.kv_cache_gid] + + def _map_draft_layers_to_groups( + self, + kv_cache_config: KVCacheConfig, + ) -> tuple[dict[str, int], dict[str, KVCacheSpec]]: + """Map Qwen draft cache owners to scheduler groups and concrete specs.""" + layer_to_gid: dict[str, int] = {} + layer_to_spec: dict[str, KVCacheSpec] = {} + for gid, group in enumerate(kv_cache_config.kv_cache_groups): + group_spec = group.kv_cache_spec + for layer_name in group.layer_names: + if layer_name not in self._draft_attn_layer_names: + continue + assert isinstance(group_spec, UniformTypeKVCacheSpecs), ( + "Qwen draft cache owners require packed KV cache groups" + ) + spec = group_spec.kv_cache_specs.get(layer_name) + assert spec is not None, ( + f"Qwen draft cache group {gid} has no spec for {layer_name}" + ) + layer_to_gid[layer_name] = gid + layer_to_spec[layer_name] = spec + + assert layer_to_spec.keys() == self._draft_attn_layer_names, ( + "Qwen draft KV cache configuration is missing layers: " + f"{sorted(self._draft_attn_layer_names - layer_to_spec.keys())}" + ) + return layer_to_gid, layer_to_spec + + +__all__ = ["Qwen4ExpMTPProposer"] diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 6a4fcbb629..d6be02effa 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -2,12 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import itertools import multiprocessing -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor from typing import TYPE_CHECKING from vllm.config import VllmConfig -from vllm.logger import init_logger from vllm.reasoning import ReasoningParserManager from vllm.tokenizers import cached_tokenizer_from_config from vllm.utils.import_utils import LazyLoader @@ -15,7 +14,6 @@ from vllm.v1.structured_output.backend_types import ( StructuredOutputBackend, StructuredOutputGrammar, - StructuredOutputOptions, ) from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend @@ -30,9 +28,6 @@ torch = LazyLoader("torch", globals(), "torch") -logger = init_logger(__name__) - - class StructuredOutputManager: """Engine-level manager for structured output requests.""" @@ -275,17 +270,63 @@ def grammar_bitmask( grammar = structured_output_request.grammar apply_bitmask = self.should_fill_bitmask(request) + reasoner = None if apply_bitmask else self._get_reasoner(request) + detect_reasoning_end = reasoner is not None + simulated_buf: list[int] | None = None + history_len = 0 + state_advancements = 0 + post_reasoning_end_in_window = False req_tokens = scheduled_spec_decode_tokens.get(req_id, ()) - for token in itertools.chain(req_tokens, (-1,)): + for i, token in enumerate(req_tokens): self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) + advance_grammar = apply_bitmask if token == -1: - # Stop advancing the grammar once we hit a padding token. apply_bitmask = False - if apply_bitmask and not grammar.is_terminated(): - accepted = grammar.accept_tokens(req_id, [token]) - assert accepted, (token, req_id, scheduled_spec_decode_tokens) - state_advancements += 1 + advance_grammar = False + elif ( + detect_reasoning_end + and reasoner is not None + and not apply_bitmask + ): + if simulated_buf is None: + history = list(request.all_token_ids) + history_len = len(history) + simulated_buf = history + list(req_tokens) + simulated = simulated_buf[: history_len + i + 1] + if reasoner.is_reasoning_end_streaming(simulated, [token]): + # The marker belongs to reasoning. Constrain the + # following draft rows, but do not advance the + # grammar through this token. + apply_bitmask = True + advance_grammar = False + post_reasoning_end_in_window = True + if advance_grammar and not grammar.is_terminated(): + if post_reasoning_end_in_window: + # These drafts were produced before their grammar + # masks existed. Validate them without logging an + # expected FSM error, then advance only valid rows. + accepted = bool(grammar.validate_tokens([token])) + if accepted: + accepted = grammar.accept_tokens(req_id, [token]) + else: + accepted = grammar.accept_tokens(req_id, [token]) + if accepted: + state_advancements += 1 + elif not post_reasoning_end_in_window: + raise AssertionError( + (token, req_id, scheduled_spec_decode_tokens) + ) + cumulative_index += 1 + + # Diffusion models do not sample a bonus position after their + # scheduled block. Autoregressive speculative decoding does. + is_diffusion = bool( + getattr(self.vllm_config.model_config, "is_diffusion", False) + ) + if not (is_diffusion and req_tokens): + bonus_apply = self.should_fill_bitmask(request) or apply_bitmask + self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),)) cumulative_index += 1 if state_advancements > 0: grammar.rollback(state_advancements) @@ -303,23 +344,31 @@ def should_fill_bitmask(self, request: "Request") -> bool: # NOTE (Hanchen) if enable_in_reasoning is True, it means that # the model needs to be constrained in reasoning. So we should always # enable the bitmask filling. + structured_req = request.structured_output_request + if self.enable_in_reasoning or ( + structured_req is not None and structured_req.reasoning_ended is True + ): + return True + reasoner = self._get_reasoner(request) if reasoner is not None: - if self.enable_in_reasoning: - return True - assert request.structured_output_request is not None - if request.structured_output_request.reasoning_ended is None: + assert structured_req is not None + if structured_req.reasoning_ended is None: # This should be removed here, but since `openai_gptoss` # is an independent code path, it is kept for now. # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, # it can be removed. - request.structured_output_request.reasoning_ended = ( - reasoner.is_reasoning_end(request.prompt_token_ids or []) + structured_req.reasoning_ended = reasoner.is_reasoning_end( + request.prompt_token_ids or [] ) - return request.structured_output_request.reasoning_ended + return structured_req.reasoning_ended return True - def should_advance(self, request: "Request") -> bool: + def should_advance( + self, + request: "Request", + new_token_ids: list[int] | None = None, + ) -> bool: if not request.use_structured_output: return False @@ -328,47 +377,74 @@ def should_advance(self, request: "Request") -> bool: if TYPE_CHECKING: assert request.structured_output_request is not None assert request.structured_output_request.grammar is not None - # by default, we should always advance - # for cases that don't use thinking mode. - reasoner = self._get_reasoner(request) - if reasoner is None: + structured_req = request.structured_output_request + if self.enable_in_reasoning or ( + structured_req is not None and structured_req.reasoning_ended is True + ): return True - # if the model needs structured in reasoning, we should advance - if self.enable_in_reasoning: + # By default, always advance requests that have no reasoning parser. + reasoner = self._get_reasoner(request) + if reasoner is None: return True - structured_req = request.structured_output_request - if structured_req.reasoning_ended: - return True + assert structured_req is not None - # Check if reasoning ends in *this* step - delta_from = request.num_computed_tokens - request.num_output_placeholders + # Use the tokens appended in this scheduler step when available. The + # placeholder-derived fallback is not reliable after partial draft + # rejection under async scheduling. all_token_ids = request.all_token_ids - start = ( - delta_from if delta_from >= 0 else max(len(all_token_ids) + delta_from, 0) - ) - if reasoner.is_reasoning_end_streaming( - all_token_ids, itertools.islice(all_token_ids, start, None) - ): + if new_token_ids: + start = len(all_token_ids) - len(new_token_ids) + delta_ids: Iterable[int] = new_token_ids + else: + delta_from = request.num_computed_tokens - request.num_output_placeholders + start = ( + delta_from + if delta_from >= 0 + else max(len(all_token_ids) + delta_from, 0) + ) + delta_ids = itertools.islice(all_token_ids, start, None) + if reasoner.is_reasoning_end_streaming(all_token_ids, delta_ids): structured_req.reasoning_ended = True - - # Reasoning just ended this step. Defer FSM advance until the next - # pass (see reasoning_ended check above) for JSON/regex/choice/grammar: - # advancing on the closing boundary token can accept tokens that still - # belong to the reasoning stream. Structural tags are the only safe - # same-step exception: they model phased output (e.g. thinking tag -> - # answer tag), and speculative decoding must run grammar.validate_tokens - # on draft tokens produced immediately after that transition. - if ( - self.vllm_config.speculative_config is not None - and structured_req.structured_output_key[0] - == StructuredOutputOptions.STRUCTURAL_TAG - ): - return True + structured_req.reasoning_end_token_index = self._find_reasoning_end_index( + reasoner, all_token_ids, start + ) + return True return False + @staticmethod + def _find_reasoning_end_index( + reasoner: "ReasoningParser", all_token_ids: Sequence[int], start: int + ) -> int: + """Return the absolute index of the final reasoning-marker token.""" + prefix = list(itertools.islice(all_token_ids, start)) + for idx in range(start, len(all_token_ids)): + token = all_token_ids[idx] + prefix.append(token) + if reasoner.is_reasoning_end_streaming(prefix, [token]): + return idx + # Multi-token markers may only be recognized from the complete delta. + # Conservatively keep the whole step out of the grammar in that case. + return len(all_token_ids) - 1 + + def trim_reasoning_for_advance( + self, request: "Request", new_token_ids: list[int] + ) -> list[int]: + """Drop reasoning content before advancing the structured grammar.""" + structured_req = request.structured_output_request + if structured_req is None: + return new_token_ids + end_idx = structured_req.reasoning_end_token_index + if end_idx is None: + return new_token_ids + first_idx = len(request.all_token_ids) - len(new_token_ids) + num_reasoning = end_idx + 1 - first_idx + if num_reasoning <= 0: + return new_token_ids + return new_token_ids[num_reasoning:] + def clear_backend(self) -> None: if self.backend is not None: self.backend.destroy() diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index a92be3d443..2fab5cc00c 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -148,11 +148,11 @@ class XgrammarGrammar(StructuredOutputGrammar): def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: """Accepts a list of tokens and advances the FSM. - Returns True if the FSM was advanced successfully. - Returns False if the FSM failed to advance. + Returns True if every constrained token was accepted. Tokens after + grammar termination are ignored. Returns False on rejection. """ if self._is_terminated: - return False + return True for token in tokens: if not self.matcher.accept_token(token): logger.error( @@ -163,7 +163,9 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: ) return False self.num_processed_tokens += 1 - self._is_terminated = self.matcher.is_terminated() + self._is_terminated = self.matcher.is_terminated() + if self._is_terminated: + break return True def validate_tokens(self, tokens: list[int]) -> list[int]: @@ -172,10 +174,15 @@ def validate_tokens(self, tokens: list[int]) -> list[int]: Returns the prefix list of tokens that are accepted by the FSM. """ + if self._is_terminated: + return [] + accepted_tokens = [] for token in tokens: if self.matcher.accept_token(token): accepted_tokens.append(token) + if self.matcher.is_terminated(): + break else: break if len(accepted_tokens) > 0: @@ -195,8 +202,9 @@ def is_terminated(self) -> bool: return self._is_terminated def reset(self): - self.num_processed_tokens = 0 self.matcher.reset() + self.num_processed_tokens = 0 + self._is_terminated = False # cf https://github.com/mlc-ai/xgrammar/blob/a32ac892676d2eedc0327416105b9b06edfb94b2/cpp/json_schema_converter.cc diff --git a/vllm/v1/structured_output/request.py b/vllm/v1/structured_output/request.py index dfa8c7efca..6caf276a11 100644 --- a/vllm/v1/structured_output/request.py +++ b/vllm/v1/structured_output/request.py @@ -23,6 +23,9 @@ class StructuredOutputRequest: params: StructuredOutputsParams _grammar: Future[StructuredOutputGrammar] | StructuredOutputGrammar | None = None reasoning_ended: bool | None = None + # Absolute index in all_token_ids of the last reasoning-marker token. + # Only the suffix after this boundary may advance the structured grammar. + reasoning_end_token_index: int | None = None reasoning_parser_kwargs: dict[str, Any] | None = None # Cached per request; do not share reasoning parsers across requests because # their behavior can depend on reasoning_parser_kwargs. diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 616fdd15bd..0e911ff560 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math + import numpy as np import torch @@ -15,6 +17,28 @@ logger = init_logger(__name__) +def get_block_table_width( + max_num_blocks: int, + block_size: int, + kernel_block_size: int | None = None, + *, + token_alignment: int | None = 128, +) -> int: + if kernel_block_size is None: + kernel_block_size = block_size + if block_size % kernel_block_size != 0: + raise ValueError( + f"kernel_block_size {kernel_block_size} must divide " + f"block_size {block_size} evenly" + ) + if token_alignment is not None: + if token_alignment <= 0: + raise ValueError("token_alignment must be positive") + alignment = token_alignment // math.gcd(token_alignment, block_size) + max_num_blocks = cdiv(max_num_blocks, alignment) * alignment + return max_num_blocks * block_size // kernel_block_size + + class BlockTable: def __init__( self, @@ -104,9 +128,7 @@ def __init__( self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size def _mark_dirty(self, row_idx: int) -> None: - self._block_table_dirty_until = max( - self._block_table_dirty_until, row_idx + 1 - ) + self._block_table_dirty_until = max(self._block_table_dirty_until, row_idx + 1) def append_row( self, diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 771f06bff5..87fcf73d1b 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -199,7 +199,9 @@ def _reshape_kv_cache( if isinstance(kv_cache_spec, AttentionSpec): has_attn = True - num_blocks_per_kv_block = kv_cache_spec.block_size // kernel_block_size + num_blocks_per_kv_block = ( + kv_cache_spec.storage_block_size // kernel_block_size + ) kernel_num_blocks = num_blocks * num_blocks_per_kv_block kv_cache_shape = group.backend.get_kv_cache_shape( kernel_num_blocks, diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 41692f58e7..d403019400 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -21,6 +21,7 @@ def __init__( cp_size: int = 1, cp_rank: int = 0, cp_interleave: int = 1, + slot_mapping_enabled: list[bool] | None = None, ): self.block_sizes = block_sizes self.kernel_block_sizes = kernel_block_sizes @@ -34,6 +35,10 @@ def __init__( self.num_kv_cache_groups = len(self.block_sizes) assert len(max_num_blocks_per_group) == self.num_kv_cache_groups + if slot_mapping_enabled is None: + slot_mapping_enabled = [True] * self.num_kv_cache_groups + assert len(slot_mapping_enabled) == self.num_kv_cache_groups + self._slot_mapping_enabled = slot_mapping_enabled self.blocks_per_kv_block = [ bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes) @@ -91,6 +96,9 @@ def init_block_table_layout_tensors(self) -> None: self.block_sizes_tensor = torch.tensor( self.kernel_block_sizes, dtype=torch.int32, device=self.device ) + self.slot_mapping_enabled = torch.tensor( + self._slot_mapping_enabled, dtype=torch.bool, device=self.device + ) self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables) def append_block_ids( @@ -158,6 +166,7 @@ def compute_slot_mappings( self.block_table_ptrs, self.block_table_strides, self.block_sizes_tensor, + self.slot_mapping_enabled, self.slot_mappings, self.slot_mappings.stride(0), self.cp_rank, @@ -229,6 +238,7 @@ def _compute_slot_mappings_kernel( block_table_ptrs, # [num_kv_cache_groups] block_table_strides, # [num_kv_cache_groups] block_sizes, # [num_kv_cache_groups] + slot_mapping_enabled, # [num_kv_cache_groups] slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] slot_mappings_stride, cp_rank, @@ -256,6 +266,7 @@ def _compute_slot_mappings_kernel( block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32) block_table_stride = tl.load(block_table_strides + group_id) block_size = tl.load(block_sizes + group_id) + mapping_enabled = tl.load(slot_mapping_enabled + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) start_idx = tl.load(query_start_loc + batch_idx) @@ -264,7 +275,9 @@ def _compute_slot_mappings_kernel( offset = i + tl.arange(0, TRITON_BLOCK_SIZE) positions = tl.load(pos + offset, mask=offset < end_idx, other=0) - block_indices = positions // (block_size * CP_SIZE) + block_indices = tl.where( + mapping_enabled, positions // (block_size * CP_SIZE), 0 + ) block_offsets = positions % (block_size * CP_SIZE) block_numbers = tl.load( block_table_ptr + req_state_idx * block_table_stride + block_indices @@ -282,6 +295,7 @@ def _compute_slot_mappings_kernel( slot_ids = block_numbers * block_size + local_offsets slot_ids = tl.where(is_local, slot_ids, PAD_ID) + slot_ids = tl.where(mapping_enabled, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index a0953c44c5..0a09cab818 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -13,6 +13,10 @@ from vllm.compilation.counter import compilation_counter from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.config.speculative import ( + get_dflash_model_draft_tokens, + uses_adaptive_dflash_lookup, +) from vllm.distributed.parallel_state import ( get_pp_group, graph_capture, @@ -34,6 +38,23 @@ logger = init_logger(__name__) +def get_explicit_cudagraph_memory_reserve(cudagraph_mode: CUDAGraphMode) -> int: + """Return an operator-provided V2 CUDA graph memory reserve in bytes.""" + reserve_mib = envs.VLLM_V2_CUDAGRAPH_MEM_MIB + if reserve_mib < 0: + raise ValueError("VLLM_V2_CUDAGRAPH_MEM_MIB must be non-negative") + if reserve_mib == 0 or cudagraph_mode == CUDAGraphMode.NONE: + return 0 + + reserve_bytes = int(reserve_mib * 1024 * 1024) + logger.info( + "Reserving %.2f GiB for V2 CUDA graphs before sizing the KV cache " + "(VLLM_V2_CUDAGRAPH_MEM_MIB)", + reserve_bytes / 2**30, + ) + return reserve_bytes + + def _use_split_sm70_mtp_cudagraphs(vllm_config: VllmConfig) -> bool: speculative_config = vllm_config.speculative_config return bool( @@ -129,6 +150,19 @@ def __init__( assert self.compilation_config is not None self.cudagraph_mode = cudagraph_mode self.decode_query_len = decode_query_len + self.decode_query_lens: tuple[int, ...] = (decode_query_len,) + speculative_config = vllm_config.speculative_config + if speculative_config is not None and uses_adaptive_dflash_lookup( + speculative_config + ): + model_query_len = 1 + get_dflash_model_draft_tokens(speculative_config) + if 1 < model_query_len < decode_query_len: + self.decode_query_lens = (model_query_len, decode_query_len) + logger.info_once( + "Capturing target decode graphs for adaptive DFlash2 " + "query lengths %s.", + self.decode_query_lens, + ) self.dp_size = vllm_config.parallel_config.data_parallel_size self.tp_size = vllm_config.parallel_config.tensor_parallel_size @@ -162,6 +196,15 @@ def __init__( self._capture_sizes = list( self.compilation_config.cudagraph_capture_sizes or [] ) + # The sizes above are normalized around the widest target query + # (q16 for LABD). Preserve the checkpoint-native q8 shape so a + # B1 short verification can still replay a graph when + # max_num_seqs=1. + if len(self.decode_query_lens) > 1: + short_query_len = min(self.decode_query_lens) + if short_query_len not in self._capture_sizes: + self._capture_sizes.append(short_query_len) + self._capture_sizes.sort() self._init_candidates() def _init_candidates(self) -> None: @@ -171,7 +214,6 @@ def _init_candidates(self) -> None: return capture_sizes = sorted(capture_sizes) - max_decode_tokens = self.max_num_reqs * self.decode_query_len decode_mode = self.cudagraph_mode.decode_mode() mixed_mode = self.cudagraph_mode.mixed_mode() separate_decode_routine = self.cudagraph_mode.separate_routine() @@ -182,19 +224,20 @@ def _init_candidates(self) -> None: for num_tokens in capture_sizes: # Capture uniform decode specfifc graphs if required # (i.e. separate decode routine) - if ( - separate_decode_routine - and decode_mode - and self.decode_query_len <= num_tokens <= max_decode_tokens - ): - desc = BatchExecutionDescriptor( - cg_mode=decode_mode, - num_tokens=num_tokens, - num_reqs=num_tokens // self.decode_query_len, - uniform_token_count=self.decode_query_len, - ) - descs_by_mode[decode_mode].append(desc) - descs_by_token_count[num_tokens].append(desc) + if separate_decode_routine and decode_mode: + for query_len in self.decode_query_lens: + if ( + query_len <= num_tokens <= self.max_num_reqs * query_len + and num_tokens % query_len == 0 + ): + desc = BatchExecutionDescriptor( + cg_mode=decode_mode, + num_tokens=num_tokens, + num_reqs=num_tokens // query_len, + uniform_token_count=query_len, + ) + descs_by_mode[decode_mode].append(desc) + descs_by_token_count[num_tokens].append(desc) if mixed_mode: # for PIECEWISE graphs there is no limit on requests when replaying diff --git a/vllm/v1/worker/gpu/mamba_align.py b/vllm/v1/worker/gpu/mamba_align.py index 52a03e83c3..c24ce798a3 100644 --- a/vllm/v1/worker/gpu/mamba_align.py +++ b/vllm/v1/worker/gpu/mamba_align.py @@ -261,8 +261,7 @@ def run_mamba_align_precopy( ) -> None: if num_reqs == 0 or not ctx.is_initialized: return - total_states = ctx.num_layers * ctx.num_state_types - grid = (num_reqs, total_states, _TEMPORAL_TILES) + grid = (num_reqs, ctx.num_states, _TEMPORAL_TILES) _precopy_mamba_align_kernel[grid]( state_idx, src_col, @@ -296,8 +295,7 @@ def run_mamba_align_postprocess( # output count while programs copying other state tensors still read it. snapshot = ctx.num_accepted_tokens_out snapshot.copy_(num_accepted_tokens) - total_states = ctx.num_layers * ctx.num_state_types - grid = (num_reqs, total_states, _TEMPORAL_TILES) + grid = (num_reqs, ctx.num_states, _TEMPORAL_TILES) _postprocess_mamba_align_kernel[grid]( snapshot, num_accepted_tokens, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 47d1da107e..ca73473d5c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -43,6 +43,7 @@ ) from vllm.model_executor.model_loader import get_model_loader from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask from vllm.utils.math_utils import cdiv @@ -50,7 +51,12 @@ from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + KVCacheConfig, + MambaSpec, + UniformTypeKVCacheSpecs, +) from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput @@ -66,6 +72,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, ModelCudaGraphManager, + get_explicit_cudagraph_memory_reserve, get_uniform_token_count, ) from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp @@ -207,6 +214,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): vocab_size=self.vocab_size, device=self.device, ) + if self.speculator is not None and hasattr(self.speculator, "set_req_states"): + self.speculator.set_req_states(self.req_states) self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, max_num_tokens=self.max_num_tokens, @@ -395,18 +404,31 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: block_sizes = [] max_num_blocks_per_group = [] + slot_mapping_enabled = [] for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + specs = tuple(spec.kv_cache_specs.values()) + assert specs + is_circular = all( + isinstance(member, CircularBufferSpec) for member in specs + ) + spec = specs[0] + else: + is_circular = isinstance(spec, CircularBufferSpec) block_sizes.append(spec.block_size) + slot_mapping_enabled.append(not is_circular) # When using DCP, each request's KV cache is sharded among different ranks. # As a result, one block on the current rank covers `block_size * cp_size` # tokens in the full, global (unsharded) sequence. - max_num_blocks = cdiv( - block_table_max_model_len, spec.block_size * self.dcp_size + max_num_blocks = ( + 1 + if is_circular + else cdiv(block_table_max_model_len, spec.block_size * self.dcp_size) ) # Align to a multiple of (128 / block_size) as required by some attention # backends such as TRTLLM (#39324) - if spec.block_size <= 128: + if not is_circular and spec.block_size <= 128: alignment = 128 // spec.block_size max_num_blocks = cdiv(max_num_blocks, alignment) * alignment # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens @@ -430,6 +452,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: cp_size=self.dcp_size, cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, + slot_mapping_enabled=slot_mapping_enabled, ) initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config @@ -499,6 +522,17 @@ def _zero_block_ids(self, block_ids: list[int]) -> None: if hasattr(self, "_kv_block_zeroer"): self._kv_block_zeroer.zero_block_ids(block_ids) + def _warmup_sm70_aux_kernels(self) -> None: + """Warm SM70 kernels whose production cache exists only in MRV2.""" + from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + _warmup_sm70_qwen_gdn_causal_conv1d, + ) + + if _warmup_sm70_qwen_gdn_causal_conv1d( + self.compilation_config.static_forward_context + ): + logger.info_once("SM70 MRV2 GDN causal-conv warmup finished.") + @torch.inference_mode() @step_eplb_after(is_dummy=True) def _dummy_run( @@ -670,7 +704,9 @@ def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int: def profile_cudagraph_memory(self) -> int: # NOTE(woosuk): It is TBD whether we keep this API or not. - return 0 + return get_explicit_cudagraph_memory_reserve( + self.compilation_config.cudagraph_mode + ) @torch.inference_mode() def capture_model(self) -> int: @@ -1025,6 +1061,29 @@ def sample( input_batch, grammar_output, ) + sm70_greedy_decode = ( + sampler_output is None + and input_batch.num_draft_tokens == 0 + and input_batch.num_reqs == 1 + and input_batch.num_tokens == 1 + and not input_batch.is_prefilling_np[0] + and grammar_output is None + and self.device.type == "cuda" + and current_platform.is_device_capability(70) + and getattr(self, "lora_config", None) is None + and hasattr(self.model, "get_top_tokens") + and self.sampler is not None + and self.sampler.can_use_sm70_greedy_token_fastpath(input_batch) + ) + if sm70_greedy_decode: + sampled = self.model.get_top_tokens(sample_hidden_states) + sampler_output = SamplerOutput( + sampled_token_ids=sampled.view(-1, 1), + logprobs_tensors=None, + num_nans=None, + num_sampled=input_batch.seq_lens.new_ones(input_batch.num_reqs), + ) + logger.info_once("SM70 MRv2 greedy TP-local pair path enabled.") if sampler_output is None: logits = self.model.compute_logits(sample_hidden_states) if grammar_output is not None: @@ -1464,7 +1523,14 @@ def sample_tokens( ), ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens - self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens) + num_draft_tokens = None + if hasattr(self.speculator, "next_num_draft_tokens"): + num_draft_tokens = self.speculator.next_num_draft_tokens() + self.draft_tokens_handler.set_draft_tokens( + input_batch, + draft_tokens, + num_draft_tokens=num_draft_tokens, + ) # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index 06b5a92c39..d7f9e6f83e 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -13,6 +13,11 @@ def init_model_state( encoder_cache: EncoderCache | None, device: torch.device, ): + get_model_state_cls = getattr(model, "get_model_state_cls", None) + if callable(get_model_state_cls): + model_state_cls = get_model_state_cls() + return model_state_cls(vllm_config, model, encoder_cache, device) + if "WhisperForConditionalGeneration" in vllm_config.model_config.architectures: from vllm.v1.worker.gpu.model_states.whisper import WhisperModelState diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 140cf97413..f38703a49c 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -31,6 +31,9 @@ prepare_dflash2_gdn_group_metadata, ) from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder +from vllm.v1.attention.backends.short_conv_attn import ( + PleShortConvAttentionMetadataBuilder, +) from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.utils import CpuGpuBuffer @@ -49,7 +52,11 @@ from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata from vllm.v1.worker.gpu.spec_decode import uses_dflash_selector_engine -from vllm.v1.worker.mamba_utils import MambaSpecDecodeGPUContext +from vllm.v1.worker.mamba_utils import ( + MambaSpecDecodeGPUContext, + get_mamba_groups, + get_mamba_types, +) from vllm.v1.worker.utils import AttentionGroup logger = init_logger(__name__) @@ -90,7 +97,11 @@ def get_extra_attn_kwargs( } if not isinstance( attn_metadata_builder, - (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder), + ( + Mamba2AttentionMetadataBuilder, + GDNAttentionMetadataBuilder, + PleShortConvAttentionMetadataBuilder, + ), ): return {} kwargs = { @@ -187,16 +198,7 @@ def _get_mamba_group_info( self, kv_cache_config: KVCacheConfig ) -> tuple[list[int], MambaSpec]: if self._mamba_spec is None: - group_ids: list[int] = [] - specs: list[MambaSpec] = [] - for group_id, group in enumerate(kv_cache_config.kv_cache_groups): - if isinstance(group.kv_cache_spec, MambaSpec): - group_ids.append(group_id) - specs.append(group.kv_cache_spec) - assert specs, "no mamba layers in the model" - assert all(specs[0] == spec for spec in specs) - self._mamba_group_ids = group_ids - self._mamba_spec = specs[0] + self._mamba_group_ids, self._mamba_spec = get_mamba_groups(kv_cache_config) return self._mamba_group_ids, self._mamba_spec def _ensure_align_ctx( @@ -205,12 +207,18 @@ def _ensure_align_ctx( mamba_group_ids: list[int], block_tables: tuple[torch.Tensor, ...], ) -> MambaSpecDecodeGPUContext: + copy_funcs = None if self._mamba_ctx is None: - copy_funcs = self.model.get_mamba_state_copy_func() + copy_funcs = self.model.get_mamba_state_copy_funcs( + get_mamba_types(kv_cache_config) + ) # This V100 closure intentionally retains the tree's default SD # layout. The official DS-row extension is unrelated to Qwen3.8's # configured path and would expand the DDTree-sensitive closure. - if get_conv_copy_spec in copy_funcs and is_conv_state_dim_first(): + if ( + any(get_conv_copy_spec in funcs for funcs in copy_funcs.values()) + and is_conv_state_dim_first() + ): raise ValueError( "MRV2 align prefix caching with speculative decoding requires " "the default SD Mamba conv-state layout on this V100 path." @@ -218,7 +226,7 @@ def _ensure_align_ctx( self._mamba_ctx = MambaSpecDecodeGPUContext.create( max_num_reqs=self.max_num_reqs, kv_cache_config=kv_cache_config, - num_state_types=len(copy_funcs), + mamba_state_copy_funcs=copy_funcs, device=self.device, make_buffer=lambda n, dtype: CpuGpuBuffer( n, @@ -229,11 +237,15 @@ def _ensure_align_ctx( ) ctx = self._mamba_ctx if not ctx.is_initialized: + if copy_funcs is None: + copy_funcs = self.model.get_mamba_state_copy_funcs( + get_mamba_types(kv_cache_config) + ) forward_context = self.vllm_config.compilation_config.static_forward_context ctx.initialize_from_forward_context( kv_cache_config, forward_context, - self.model.get_mamba_state_copy_func(), + copy_funcs, [block_tables[group_id] for group_id in mamba_group_ids], ) return ctx diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 8bf884fd9b..7ae94c43c2 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -62,6 +62,24 @@ def apply_staged_writes(self) -> None: self.bad_words_state.apply_staged_writes() self.logprob_token_ids_state.apply_staged_writes() + def can_use_sm70_greedy_token_fastpath(self, input_batch: InputBatch) -> bool: + """Whether TP-local argmax is equivalent to the full MRv2 sampler.""" + if not envs.VLLM_SM70_GREEDY_TOKEN_FASTPATH or self.compute_nans: + return False + + req_indices = input_batch.idx_mapping_np + if not np.all(self.sampling_states.temperature.np[req_indices] == 0.0): + return False + if self.sampling_states.max_num_logprobs(req_indices) != NO_LOGPROBS: + return False + if self.logprob_token_ids_state.max_num_token_ids(req_indices) != 0: + return False + if np.any(self.logit_bias_state.use_logit_bias[req_indices]): + return False + if np.any(self.penalties_state.use_penalty[req_indices]): + return False + return not np.any(self.bad_words_state.num_bad_words.np[req_indices]) + def __call__( self, logits: torch.Tensor, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 9d08b77ee5..d645c2a420 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -9,6 +9,7 @@ from vllm.config import VllmConfig, replace from vllm.config.compilation import CUDAGraphMode +from vllm.config.speculative import get_dflash_model_draft_tokens from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.triton_utils import tl, triton @@ -53,8 +54,19 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # Multimodal inputs not currently supported. self.supports_mm_inputs = False - # Each request emits exactly (bonus + N mask) query tokens per step. - self.num_query_per_req = 1 + self.num_speculative_steps + # A lookup-augmented DFlash2 proposal may be wider than the checkpoint + # block. The model still emits only its trained block; lookup fills the + # remaining target-verification positions. + self.draft_block = get_dflash_model_draft_tokens(self.speculative_config) + self.num_query_per_req = 1 + self.draft_block + if self.draft_block < self.num_speculative_steps: + logger.info( + "%s emits %d model drafts and lookup may fill %d additional " + "target-verification positions.", + self._speculator_name, + self.draft_block, + self.num_speculative_steps - self.draft_block, + ) self.parallel_drafting_token_id = get_parallel_drafting_token_id( self.draft_model_config.hf_config @@ -88,7 +100,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ) # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens). - max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps + max_num_sampled_tokens = self.max_num_reqs * self.draft_block self.sample_indices = torch.zeros( max_num_sampled_tokens, dtype=torch.int64, device=device ) @@ -104,7 +116,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # [0, 1, ..., N-1, 0, 1, ..., N-1, ...] -> the per-token column index into # draft_logits[req, step, :]. self.sample_col = torch.arange( - self.num_speculative_steps, dtype=torch.int32, device=device + self.draft_block, dtype=torch.int32, device=device ).repeat(self.max_num_reqs) self.query_cudagraph_manager: DFlashCudaGraphManager | None = None @@ -271,7 +283,7 @@ def _generate_draft( num_tokens_across_dp, cudagraph_runtime_mode, ) - num_sample = num_reqs * self.num_speculative_steps + num_sample = num_reqs * self.draft_block sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]] # sample_pos is the predicted token's position Q; verification keys # Gumbel by the predecessor (Q-1). sample_draft adds +1, so pass Q-2. @@ -284,8 +296,8 @@ def _generate_draft( self.sample_col[:num_sample], self.draft_logits, ) - self.draft_tokens[:num_reqs] = draft_tokens.view( - num_reqs, self.num_speculative_steps + self.draft_tokens[:num_reqs, : self.draft_block] = draft_tokens.view( + num_reqs, self.draft_block ) def _build_draft_attn_metadata( @@ -336,6 +348,14 @@ def _prepare_ngram_assist( """Prepare an optional draftless proposal; return whether it is complete.""" return False + def _prepare_proposal_runtime( + self, + input_batch: InputBatch, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + ) -> None: + """Refresh optional proposal-controller state before graph replay.""" + def _apply_ngram_assist(self, num_reqs: int) -> None: """Override model proposals for ngram-hit rows, if configured.""" @@ -378,6 +398,7 @@ def propose( self.draft_max_seq_len = min( max_seq_len + self.num_query_per_req, self.max_model_len ) + self._prepare_proposal_runtime(input_batch, num_sampled, num_rejected) # NOTE: To avoid CPU-GPU synchronization without CPU knowing the # number of rejected tokens, we maintain the size of input_ids and @@ -443,7 +464,7 @@ def propose( self.block_tables.cp_interleave, self.parallel_drafting_token_id, self.num_query_per_req, - self.num_speculative_steps, + self.draft_block, self.max_num_reqs, self.max_num_tokens, self.max_model_len, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py b/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py new file mode 100644 index 0000000000..0bc065907f --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py @@ -0,0 +1,326 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPU lookup-augmented block drafting for DFlash2. + +Adapted for MRV2/SM70 from syv-ai/qwen38-27b-rtx3090's Apache-2.0 +``dflash2-lookup-drafting.patch`` (fixed source revision is documented in the +design note). +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, triton + +_SCORE_STRIDE = tl.constexpr(1 << 32) + + +@triton.jit +def _suffix_lookup_kernel( + all_token_ids_ptr, + all_token_ids_stride, + total_len_ptr, + idx_mapping_ptr, + idx_mapping_stride, + eligible_ptr, + out_tokens_ptr, + out_tokens_stride, + out_len_ptr, + out_valid_ptr, + search_max, + k: tl.constexpr, + NMAX: tl.constexpr, + NMIN: tl.constexpr, + BLOCK: tl.constexpr, + BLOCK_K: tl.constexpr, +): + req = tl.program_id(0) + tl.store(out_len_ptr + req, 0) + tl.store(out_valid_ptr + req, 0) + if tl.load(eligible_ptr + req) == 0: + return + req_state = tl.load(idx_mapping_ptr + req * idx_mapping_stride) + if req_state < 0: + return + + total_len = tl.load(total_len_ptr + req_state) + if total_len < NMIN + 2: + return + base = all_token_ids_ptr + req_state.to(tl.int64) * all_token_ids_stride + suffix_end = total_len - 1 + search_end = suffix_end + search_start = tl.maximum(NMIN - 1, total_len - search_max) + best = tl.full([], -1, tl.int64) + + # Select the longest suffix match and then the most recent occurrence. + # Candidate matches may overlap the current suffix, which preserves + # periodic patterns such as indentation, list markers, and code fences. + for start in range(search_start, search_end, BLOCK): + candidate_end = start + tl.arange(0, BLOCK) + alive = (candidate_end >= search_start) & (candidate_end < search_end) + for offset in range(NMIN): + suffix_token = tl.load(base + suffix_end - offset) + candidate_token = tl.load( + base + candidate_end - offset, + mask=alive & ((candidate_end - offset) >= 0), + other=-1, + ) + alive &= candidate_token == suffix_token + + if tl.max(alive.to(tl.int32)) > 0: + match_len = tl.where(alive, NMIN, 0) + longer = alive + for offset in range(NMIN, NMAX): + suffix_token = tl.load( + base + suffix_end - offset, + mask=(suffix_end - offset) >= 0, + other=-2, + ) + candidate_token = tl.load( + base + candidate_end - offset, + mask=longer & ((candidate_end - offset) >= 0), + other=-1, + ) + longer &= candidate_token == suffix_token + match_len = tl.where(longer, match_len + 1, match_len) + score = tl.where( + alive, + match_len.to(tl.int64) * _SCORE_STRIDE + candidate_end.to(tl.int64), + -1, + ) + best = tl.maximum(best, tl.max(score, axis=0)) + + if best < 0: + return + match_len = (best // _SCORE_STRIDE).to(tl.int32) + candidate_end = (best % _SCORE_STRIDE).to(tl.int32) + valid = tl.minimum(k, suffix_end - candidate_end) + offsets = tl.arange(0, BLOCK_K) + tokens = tl.load( + base + candidate_end + 1 + offsets, + mask=offsets < valid, + other=0, + ) + tl.store( + out_tokens_ptr + req * out_tokens_stride + offsets, + tokens, + mask=offsets < k, + ) + tl.store(out_len_ptr + req, match_len) + tl.store(out_valid_ptr + req, valid) + + +@triton.jit +def _fuse_draft_kernel( + draft_tokens_ptr, + draft_stride, + lookup_tokens_ptr, + lookup_stride, + match_len_ptr, + valid_ptr, + use_ptr, + idx_mapping_ptr, + idx_mapping_stride, + hits_ptr, + take_flags_ptr, + nmin, + nstrong, + agree_min, + nmin_tail, + long_min, + draft_block, + k: tl.constexpr, + BLOCK_K: tl.constexpr, +): + req = tl.program_id(0) + tl.store(take_flags_ptr + req, 0) + if tl.load(idx_mapping_ptr + req * idx_mapping_stride) < 0: + return + + offsets = tl.arange(0, BLOCK_K) + mask = offsets < k + match_len = tl.load(match_len_ptr + req) + valid = tl.load(valid_ptr + req) + drafted = tl.load( + draft_tokens_ptr + req * draft_stride + offsets, + mask=mask, + other=0, + ) + looked_up = tl.load( + lookup_tokens_ptr + req * lookup_stride + offsets, + mask=mask, + other=0, + ).to(tl.int64) + + disagreement = ( + (drafted != looked_up) & (offsets < valid) & (offsets < draft_block) & mask + ) + agreement = tl.minimum(tl.min(tl.where(disagreement, offsets, draft_block)), valid) + take_head = (match_len >= nstrong) | ( + (match_len >= nmin) & (agreement >= agree_min) + ) + tail = offsets >= draft_block + take_tail = (match_len >= nmin_tail) & (take_head | (agreement >= draft_block)) + from_lookup = tl.where(tail, take_tail, take_head) & (offsets < valid) & mask + tl.store( + draft_tokens_ptr + req * draft_stride + offsets, + looked_up, + mask=from_lookup, + ) + + # Every tail row needs an explicit point-mass proposal whenever a long + # block is replayed, including sticky transition steps without a match. + use = (from_lookup | tail) & mask + tl.store(use_ptr + req * k + offsets, use.to(tl.int32), mask=mask) + has_tail = (match_len >= long_min) & (valid > draft_block) & take_tail + tl.store(take_flags_ptr + req, has_tail.to(tl.int32)) + tl.atomic_add(hits_ptr, take_head.to(tl.int64)) + + +@triton.jit +def _point_mass_draft_logits_kernel( + draft_logits_ptr, + cached_candidate_ptr, + cached_score_ptr, + draft_tokens_ptr, + draft_tokens_stride, + use_ptr, + idx_mapping_ptr, + idx_mapping_stride, + cache_stride_0, + cache_stride_1, + draft_logits_stride_0, + draft_logits_stride_1, + num_steps: tl.constexpr, + top_k: tl.constexpr, + BLOCK_K: tl.constexpr, + CACHE_SCORES: tl.constexpr, +): + flat = tl.program_id(0) + req = flat // num_steps + step = flat % num_steps + req_state = tl.load(idx_mapping_ptr + req * idx_mapping_stride) + if req_state < 0 or tl.load(use_ptr + flat) == 0: + return + + offsets = tl.arange(0, BLOCK_K) + mask = offsets < top_k + cache_base = req_state * cache_stride_0 + step * cache_stride_1 + old_ids = tl.load(cached_candidate_ptr + cache_base + offsets, mask=mask) + logits_base = ( + draft_logits_ptr + + req_state * draft_logits_stride_0 + + step * draft_logits_stride_1 + ) + tl.store(logits_base + old_ids, -float("inf"), mask=mask) + + token = tl.load(draft_tokens_ptr + req * draft_tokens_stride + step) + is_proposal = offsets == 0 + new_ids = tl.where(is_proposal, token, 0) + new_scores = tl.where(is_proposal, 0.0, -float("inf")) + tl.store(cached_candidate_ptr + cache_base + offsets, new_ids, mask=mask) + if CACHE_SCORES: + tl.store(cached_score_ptr + cache_base + offsets, new_scores, mask=mask) + tl.store(logits_base + token, 0.0) + + +def suffix_lookup( + all_token_ids: torch.Tensor, + total_len: torch.Tensor, + idx_mapping: torch.Tensor, + eligible: torch.Tensor, + num_reqs: int, + num_draft_tokens: int, + *, + idx_mapping_stride: int = 1, + nmax: int = 12, + nmin: int = 6, + search_max: int = 1 << 30, + out_tokens: torch.Tensor | None = None, + out_len: torch.Tensor | None = None, + out_valid: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Find recent continuations of each request's current suffix.""" + if not 1 <= nmin <= nmax: + raise ValueError("lookup ngram bounds must satisfy 1 <= nmin <= nmax") + if num_draft_tokens <= 0: + raise ValueError("num_draft_tokens must be positive") + device = idx_mapping.device + if out_tokens is None: + out_tokens = torch.zeros( + (num_reqs, num_draft_tokens), dtype=torch.int32, device=device + ) + if out_len is None: + out_len = torch.zeros(num_reqs, dtype=torch.int32, device=device) + if out_valid is None: + out_valid = torch.zeros(num_reqs, dtype=torch.int32, device=device) + _suffix_lookup_kernel[(num_reqs,)]( + all_token_ids, + all_token_ids.stride(0), + total_len, + idx_mapping, + idx_mapping_stride, + eligible, + out_tokens, + out_tokens.stride(0), + out_len, + out_valid, + search_max, + k=num_draft_tokens, + NMAX=nmax, + NMIN=nmin, + BLOCK=1024, + BLOCK_K=triton.next_power_of_2(num_draft_tokens), + num_warps=4, + ) + return out_tokens, out_len, out_valid + + +def fuse_draft( + draft_tokens: torch.Tensor, + lookup_tokens: torch.Tensor, + match_len: torch.Tensor, + valid: torch.Tensor, + use: torch.Tensor, + idx_mapping: torch.Tensor, + hits: torch.Tensor, + num_reqs: int, + num_draft_tokens: int, + *, + draft_block: int, + idx_mapping_stride: int = 1, + nmin: int = 6, + nstrong: int = 6, + agree_min: int = 0, + nmin_tail: int = 4, + long_min: int = 6, + take_flags: torch.Tensor | None = None, +) -> None: + """Fuse lookup continuations into a model-drafted prefix in place.""" + if take_flags is None: + take_flags = torch.zeros( + num_reqs, dtype=torch.int32, device=draft_tokens.device + ) + _fuse_draft_kernel[(num_reqs,)]( + draft_tokens, + draft_tokens.stride(0), + lookup_tokens, + lookup_tokens.stride(0), + match_len, + valid, + use, + idx_mapping, + idx_mapping_stride, + hits, + take_flags, + nmin, + nstrong, + agree_min, + nmin_tail, + long_min, + draft_block, + k=num_draft_tokens, + BLOCK_K=triton.next_power_of_2(num_draft_tokens), + num_warps=1, + ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index 64f987238a..9de53db15e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -13,6 +13,11 @@ from vllm.triton_utils import tl, triton from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator +from vllm.v1.worker.gpu.spec_decode.dflash2.lookup import ( + _point_mass_draft_logits_kernel, + fuse_draft, + suffix_lookup, +) from vllm.v1.worker.gpu.spec_decode.dflash2.ngram_assist import DFlash2NgramAssist logger = init_logger(__name__) @@ -177,6 +182,7 @@ def _cache_draft_logits_kernel( draft_logits_stride_0, draft_logits_stride_1, num_steps: tl.constexpr, + cache_steps: tl.constexpr, top_k: tl.constexpr, BLOCK_K: tl.constexpr, CACHE_SCORES: tl.constexpr, @@ -187,7 +193,7 @@ def _cache_draft_logits_kernel( offsets = tl.arange(0, BLOCK_K) mask = (req_state >= 0) & (offsets < top_k) candidate_base = flat * top_k - cache_base = (req_state * num_steps + step) * top_k + cache_base = (req_state * cache_steps + step) * top_k old_token_ids = tl.load(cached_candidate_ptr + cache_base + offsets, mask=mask) logits_base = ( draft_logits_ptr @@ -261,6 +267,58 @@ def _apply_ngram_draft_kernel( tl.store(logits_base + token, 0.0, mask=valid) +@triton.jit +def _prepare_lookup_controller_flags_kernel( + take_flags_ptr, + emitted_ptr, + out_ptr, + full_emitted, + num_reqs, + BLOCK: tl.constexpr, +): + offsets = tl.arange(0, BLOCK) + mask = offsets < num_reqs + take = tl.load(take_flags_ptr + offsets, mask=mask, other=0) > 0 + emitted = tl.load(emitted_ptr + offsets, mask=mask, other=0) + wants_long = take & (emitted >= full_emitted) + tl.store(out_ptr + offsets, wants_long.to(tl.int32), mask=mask) + + +def _advance_lookup_controller( + *, + want: bool | None, + num_reqs: int, + entry_streak: int, + sticky_steps: int, + last_want: bool, + want_streak: int, + sticky_remaining: int, + long_active: bool, +) -> tuple[bool, int, int, bool]: + """Advance the host-only adaptive q8/q16 controller. + + ``want=None`` means the asynchronous flag copy has not landed. Preserve + the prior decision in that case instead of synchronizing the decode + stream merely to choose a verification width. + """ + if want is None: + return last_want, want_streak, sticky_remaining, long_active + + if want: + want_streak = want_streak + 1 if last_want else 1 + if want_streak >= max(entry_streak, 1): + long_active = True + sticky_remaining = max(sticky_steps, 0) if num_reqs == 1 else 0 + else: + want_streak = 0 + if num_reqs == 1 and long_active and sticky_remaining > 0: + sticky_remaining -= 1 + else: + long_active = False + sticky_remaining = 0 + return want, want_streak, sticky_remaining, long_active + + class DFlash2Speculator(DFlashSpeculator): _speculator_name = "DFlash2" @@ -272,22 +330,37 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): torch.arange(self.max_num_reqs, dtype=torch.int64, device=device) * self.num_query_per_req ) + self._selector_tokens = torch.empty( + self.max_num_reqs, + self.draft_block, + dtype=self.draft_tokens.dtype, + device=device, + ) self._selector_scores = torch.empty( self.max_num_reqs, - self.num_speculative_steps, + self.draft_block, self.selector_top_k, dtype=torch.float32, device=device, ) self._cached_candidate_ids = torch.zeros( - self._selector_scores.shape, dtype=torch.int64, device=device + self.max_num_reqs, + self.num_speculative_steps, + self.selector_top_k, + dtype=torch.int64, + device=device, ) self._cached_candidate_scores = None if ( self.draft_logits is not None and envs.VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION ): - self._cached_candidate_scores = torch.empty_like(self._selector_scores) + self._cached_candidate_scores = torch.full( + self._cached_candidate_ids.shape, + -float("inf"), + dtype=torch.float32, + device=device, + ) self._selector_path_state = torch.empty( self.max_num_reqs, dtype=torch.int32, device=device ) @@ -300,7 +373,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ): packed_shape = ( self.max_num_reqs, - self.num_speculative_steps, + self.draft_block, self.selector_top_k, ) self._alignment_candidate_ids = torch.empty( @@ -314,7 +387,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): dtype=torch.float32, device=device, ) - self._use_sm70_tail = _requires_sm70_tail(device, self.num_speculative_steps) + self._use_sm70_tail = _requires_sm70_tail(device, self.draft_block) if self.draft_logits is not None: # The cache kernel writes only K columns; all other vocabulary # columns must remain impossible. @@ -325,8 +398,14 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self._ngram_rounds = 0 self._ngram_skipped_rounds = 0 speculative_config = getattr(self, "speculative_config", None) - if speculative_config is not None and getattr( - speculative_config, "ngram_assist", False + ngram_assist = bool( + speculative_config is not None + and getattr(speculative_config, "ngram_assist", False) + ) + if ( + speculative_config is not None + and ngram_assist + and self.draft_block == self.num_speculative_steps ): min_ngram = speculative_config.prompt_lookup_min max_ngram = speculative_config.prompt_lookup_max @@ -369,11 +448,229 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.num_speculative_steps, ) + self._lookup_enabled = bool( + ngram_assist and self.draft_block < self.num_speculative_steps + ) + self._req_states = None + self._lookup_current_req_key: tuple[int, ...] = () + self._lookup_current_eligible = False + self._lookup_last_emitted: torch.Tensor | None = None + if self._lookup_enabled: + if device.type != "cuda": + raise ValueError("Lookup-augmented DFlash2 requires a CUDA device") + assert speculative_config is not None + assert speculative_config.prompt_lookup_min is not None + assert speculative_config.prompt_lookup_max is not None + self._lookup_nmin = int(speculative_config.prompt_lookup_min) + self._lookup_nmax = int(speculative_config.prompt_lookup_max) + self._lookup_nstrong = envs.VLLM_DFLASH2_LOOKUP_NSTRONG + self._lookup_agree = envs.VLLM_DFLASH2_LOOKUP_AGREE + self._lookup_nmin_tail = envs.VLLM_DFLASH2_LOOKUP_NMIN_TAIL + self._lookup_long_min = envs.VLLM_DFLASH2_LOOKUP_LONG_MIN + self._lookup_search = envs.VLLM_DFLASH2_LOOKUP_SEARCH + self._lookup_adaptive = envs.VLLM_DFLASH2_LOOKUP_ADAPTIVE + self._lookup_entry_streak = envs.VLLM_DFLASH2_LOOKUP_ENTRY_STREAK + self._lookup_sticky_steps = envs.VLLM_DFLASH2_LOOKUP_STICKY + self._lookup_cheap_context = envs.VLLM_DFLASH2_LOOKUP_CHEAP_CONTEXT + self._lookup_tokens = torch.zeros( + self.max_num_reqs, + self.num_speculative_steps, + dtype=torch.int32, + device=device, + ) + self._lookup_match_len = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + self._lookup_valid = torch.zeros_like(self._lookup_match_len) + self._lookup_eligible = torch.zeros_like(self._lookup_match_len) + self._lookup_use = torch.zeros( + self.max_num_reqs, + self.num_speculative_steps, + dtype=torch.int32, + device=device, + ) + self._lookup_take_flags = torch.zeros_like(self._lookup_match_len) + self._lookup_controller_flags = torch.zeros_like(self._lookup_match_len) + self._lookup_hits = torch.zeros((), dtype=torch.int64, device=device) + self._lookup_copy_stream = torch.cuda.Stream(device=device) + self._lookup_copy_event = torch.cuda.Event() + self._lookup_flags_cpu = torch.zeros( + self.max_num_reqs, + dtype=torch.int32, + device="cpu", + pin_memory=True, + ) + self._lookup_copy_pending = False + self._lookup_pending_req_key: tuple[int, ...] = () + self._lookup_pending_num_reqs = 0 + self._lookup_last_want = False + self._lookup_want_streak = 0 + self._lookup_sticky_remaining = 0 + self._lookup_long_active = False + self._lookup_controller_req_key: tuple[int, ...] = () + self._lookup_last_verify_tokens = 0 + self._lookup_q8_rounds = 0 + self._lookup_q16_rounds = 0 + logger.info( + "Enabled GPU lookup-augmented DFlash2: model drafts=%d, " + "max target drafts=%d, ngram=[%d,%d], adaptive=%s.", + self.draft_block, + self.num_speculative_steps, + self._lookup_nmin, + self._lookup_nmax, + self._lookup_adaptive, + ) + @property def requires_host_token_state(self) -> bool: """Whether the runner must expose async samples and request history.""" return self._ngram_assist is not None + def set_req_states(self, req_states) -> None: + """Expose the request token history used by the device lookup.""" + self._req_states = req_states + + def _reset_lookup_controller(self, req_key: tuple[int, ...]) -> None: + self._lookup_controller_req_key = req_key + self._lookup_last_want = False + self._lookup_want_streak = 0 + self._lookup_sticky_remaining = 0 + self._lookup_long_active = False + + def _record_lookup_width(self, draft_tokens: int, reason: str) -> int: + """Record the asynchronously selected target verification width.""" + verify_tokens = 1 + draft_tokens + if draft_tokens == self.draft_block: + self._lookup_q8_rounds += 1 + else: + self._lookup_q16_rounds += 1 + if ( + envs.VLLM_DFLASH_PROFILE + and verify_tokens != self._lookup_last_verify_tokens + ): + logger.info( + "DFlash2 lookup target verifier selected q%d (%s); " + "q8_rounds=%d, q16_rounds=%d.", + verify_tokens, + reason, + self._lookup_q8_rounds, + self._lookup_q16_rounds, + ) + self._lookup_last_verify_tokens = verify_tokens + return draft_tokens + + def _prepare_proposal_runtime( + self, + input_batch, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + ) -> None: + """Refresh lookup eligibility and controller inputs before replay.""" + del num_rejected + if not self._lookup_enabled: + return + + num_reqs = input_batch.num_reqs + req_key = tuple(int(index) for index in input_batch.idx_mapping_np[:num_reqs]) + self._lookup_current_req_key = req_key + self._lookup_last_emitted = num_sampled[:num_reqs] + + # Grammar validation keeps the checkpoint-native q8 scheduler + # contract. Prefill rows have no stable generated suffix to extend. + self._lookup_eligible.zero_() + eligible = np.logical_not(input_batch.is_prefilling_np[:num_reqs]) + if input_batch.has_structured_output_reqs: + eligible.fill(False) + self._lookup_current_eligible = bool(num_reqs and eligible.all()) + eligible_cpu = torch.from_numpy(eligible.astype(np.int32, copy=False)) + self._lookup_eligible[:num_reqs].copy_(eligible_cpu, non_blocking=True) + + if req_key != self._lookup_controller_req_key: + self._reset_lookup_controller(req_key) + + def _consume_lookup_flags(self) -> bool | None: + """Return the last completed batch-wide lookup decision, if ready.""" + if not self._lookup_copy_pending or not self._lookup_copy_event.query(): + return None + + pending_key = self._lookup_pending_req_key + pending_num_reqs = self._lookup_pending_num_reqs + self._lookup_copy_pending = False + if pending_key != self._lookup_current_req_key or pending_num_reqs <= 0: + return False + return bool(self._lookup_flags_cpu[:pending_num_reqs].numpy().all()) + + def _queue_lookup_flags(self, num_reqs: int) -> None: + """Asynchronously copy the current q16 signal into pinned host memory.""" + if self._lookup_copy_pending or self._lookup_last_emitted is None: + return + + block = triton.next_power_of_2(max(num_reqs, 1)) + _prepare_lookup_controller_flags_kernel[(1,)]( + self._lookup_take_flags, + self._lookup_last_emitted, + self._lookup_controller_flags, + 1 + self.draft_block, + num_reqs, + BLOCK=block, + num_warps=1, + ) + current_stream = torch.cuda.current_stream(self.device) + self._lookup_copy_stream.wait_stream(current_stream) + with torch.cuda.stream(self._lookup_copy_stream): + self._lookup_flags_cpu[:num_reqs].copy_( + self._lookup_controller_flags[:num_reqs], non_blocking=True + ) + self._lookup_copy_event.record(self._lookup_copy_stream) + self._lookup_pending_req_key = self._lookup_current_req_key + self._lookup_pending_num_reqs = num_reqs + self._lookup_copy_pending = True + + def next_num_draft_tokens(self) -> int: + """Choose q8 or q16 for the next target verification step.""" + if not self._lookup_enabled: + return self.num_speculative_steps + + num_reqs = len(self._lookup_current_req_key) + if num_reqs == 0 or not self._lookup_current_eligible: + self._reset_lookup_controller(self._lookup_current_req_key) + return self._record_lookup_width(self.draft_block, "ineligible") + + want = self._consume_lookup_flags() + self._queue_lookup_flags(num_reqs) + if not self._lookup_adaptive: + return self._record_lookup_width( + self.num_speculative_steps, "adaptive-disabled" + ) + if ( + self._lookup_cheap_context > 0 + and self.draft_max_seq_len <= self._lookup_cheap_context + ): + return self._record_lookup_width( + self.num_speculative_steps, "cheap-context" + ) + + ( + self._lookup_last_want, + self._lookup_want_streak, + self._lookup_sticky_remaining, + self._lookup_long_active, + ) = _advance_lookup_controller( + want=want, + num_reqs=num_reqs, + entry_streak=self._lookup_entry_streak, + sticky_steps=self._lookup_sticky_steps, + last_want=self._lookup_last_want, + want_streak=self._lookup_want_streak, + sticky_remaining=self._lookup_sticky_remaining, + long_active=self._lookup_long_active, + ) + width = ( + self.num_speculative_steps if self._lookup_long_active else self.draft_block + ) + reason = "strong-copy" if self._lookup_long_active else "adaptive-default" + return self._record_lookup_width(width, reason) + def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]: # The selector walk and rejection sampler must consume identical scores. # BF16 rounding measurably changes candidate order, so keep this FP32. @@ -391,11 +688,7 @@ def _sample_path( scores = scores.contiguous() candidate_ids = candidate_ids.contiguous() block_k = triton.next_power_of_2(self.selector_top_k) - walk_steps = ( - self.num_speculative_steps - 1 - if self._use_sm70_tail - else self.num_speculative_steps - ) + walk_steps = self.draft_block - 1 if self._use_sm70_tail else self.draft_block _selector_walk_kernel[(num_reqs,)]( scores, candidate_ids, @@ -403,10 +696,10 @@ def _sample_path( self.sample_idx_mapping, self.temperature, self.seeds, - self.draft_tokens, + self._selector_tokens, self._selector_scores, self._selector_path_state, - num_steps=self.num_speculative_steps, + num_steps=self.draft_block, walk_steps=walk_steps, top_k=self.selector_top_k, BLOCK_K=block_k, @@ -422,10 +715,10 @@ def _sample_path( self.sample_idx_mapping, self.temperature, self.seeds, - self.draft_tokens, + self._selector_tokens, self._selector_scores, self._selector_path_state, - num_steps=self.num_speculative_steps, + num_steps=self.draft_block, top_k=self.selector_top_k, BLOCK_K=block_k, SAMPLE_PROBABILISTIC=self.draft_logits is not None, @@ -447,7 +740,8 @@ def _cache_draft_logits(self, candidate_ids: torch.Tensor, num_sample: int) -> N self.sample_idx_mapping, draft_logits.stride(0), draft_logits.stride(1), - num_steps=self.num_speculative_steps, + num_steps=self.draft_block, + cache_steps=self.num_speculative_steps, top_k=self.selector_top_k, BLOCK_K=block_k, CACHE_SCORES=cached_scores is not None, @@ -576,6 +870,71 @@ def _apply_ngram_assist(self, num_reqs: int) -> None: num_warps=1, ) + def _apply_lookup(self, num_reqs: int) -> None: + """Fuse a history continuation into the DFlash2 proposal exactly.""" + if not self._lookup_enabled or self._req_states is None: + return + + tokens, match_len, valid = suffix_lookup( + self._req_states.all_token_ids.gpu, + self._req_states.total_len.gpu, + self.sample_idx_mapping, + self._lookup_eligible, + num_reqs, + self.num_speculative_steps, + idx_mapping_stride=self.draft_block, + nmax=self._lookup_nmax, + nmin=self._lookup_nmin, + search_max=self._lookup_search, + out_tokens=self._lookup_tokens, + out_len=self._lookup_match_len, + out_valid=self._lookup_valid, + ) + fuse_draft( + self.draft_tokens, + tokens, + match_len, + valid, + self._lookup_use, + self.sample_idx_mapping, + self._lookup_hits, + num_reqs, + self.num_speculative_steps, + draft_block=self.draft_block, + idx_mapping_stride=self.draft_block, + nmin=self._lookup_nmin, + nstrong=self._lookup_nstrong, + agree_min=self._lookup_agree, + nmin_tail=self._lookup_nmin_tail, + long_min=self._lookup_long_min, + take_flags=self._lookup_take_flags, + ) + + draft_logits = self.draft_logits + if draft_logits is None: + return + cached_scores = self._cached_candidate_scores + block_k = triton.next_power_of_2(self.selector_top_k) + _point_mass_draft_logits_kernel[(num_reqs * self.num_speculative_steps,)]( + draft_logits, + self._cached_candidate_ids, + self._selector_scores if cached_scores is None else cached_scores, + self.draft_tokens, + self.draft_tokens.stride(0), + self._lookup_use, + self.sample_idx_mapping, + self.draft_block, + self._cached_candidate_ids.stride(0), + self._cached_candidate_ids.stride(1), + draft_logits.stride(0), + draft_logits.stride(1), + num_steps=self.num_speculative_steps, + top_k=self.selector_top_k, + BLOCK_K=block_k, + CACHE_SCORES=cached_scores is not None, + num_warps=1, + ) + def _generate_draft( self, num_reqs: int, @@ -592,15 +951,15 @@ def _generate_draft( num_tokens_across_dp, cudagraph_runtime_mode, ) - num_sample = num_reqs * self.num_speculative_steps + num_sample = num_reqs * self.draft_block hidden_states = last_hidden_states[self.sample_indices[:num_sample]].view( - num_reqs, self.num_speculative_steps, -1 + num_reqs, self.draft_block, -1 ) candidate_ids, unary_logits = self.model.compute_candidates( hidden_states.flatten(0, 1) ) candidate_ids = candidate_ids.view( - num_reqs, self.num_speculative_steps, self.selector_top_k + num_reqs, self.draft_block, self.selector_top_k ) unary_logits = unary_logits.view_as(candidate_ids) anchor_token_ids = self.input_buffers.input_ids[self._anchor_indices[:num_reqs]] @@ -617,5 +976,11 @@ def _generate_draft( self._alignment_unary_logits[:num_reqs].copy_(unary_logits) self._alignment_lattice_scores[:num_reqs].copy_(scores) self._sample_path(candidate_ids, scores, num_reqs) + self.draft_tokens[:num_reqs, : self.draft_block].copy_( + self._selector_tokens[:num_reqs] + ) + if self.draft_block < self.num_speculative_steps: + self.draft_tokens[:num_reqs, self.draft_block :].zero_() if self.draft_logits is not None: self._cache_draft_logits(candidate_ids, num_sample) + self._apply_lookup(num_reqs) diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 8376a51558..c09186ceba 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -6,7 +6,7 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config import SpeculativeConfig, VllmConfig, get_layers_from_vllm_config from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger @@ -43,8 +43,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config self.device = device - self.speculative_config = vllm_config.speculative_config - assert self.speculative_config is not None + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + self.speculative_config: SpeculativeConfig = speculative_config self.method = self.speculative_config.method self.num_speculative_steps = self.speculative_config.num_speculative_tokens self.draft_model_config = self.speculative_config.draft_model_config @@ -65,6 +66,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vocab_size = self.draft_model_config.get_vocab_size() self.dtype = vllm_config.model_config.dtype self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel + self.use_local_argmax_reduction = ( + self.speculative_config.use_local_argmax_reduction + ) + self.share_mtp_topk_indices = False # DP configuration self.dp_size = vllm_config.parallel_config.data_parallel_size @@ -157,6 +162,17 @@ def load_model(self, target_model: nn.Module) -> None: ).keys() self.model = load_eagle_model(target_model, self.vllm_config) + self._validate_local_argmax_reduction() + + draft_hf_config = self.draft_model_config.hf_config + self.share_mtp_topk_indices = ( + self.method == "mtp" + and getattr(draft_hf_config, "index_share_for_mtp_iteration", False) + and hasattr(self.model.model, "set_skip_topk") + and hasattr(self.model.model, "compact_topk_indices") + ) + if self.share_mtp_topk_indices: + logger.info("Reusing target-aligned step-0 QSA indices for MTP steps 1+.") all_attn_layers = get_layers_from_vllm_config( self.vllm_config, @@ -222,22 +238,44 @@ def run_model( hidden_states=self.hidden_states[:num_tokens], inputs_embeds=inputs_embeds, ) - if self.method == "mtp": + # Some MTP models declare a single-tensor contract but return + # (logits_hidden, feedback_hidden) for final-norm correctness. + if isinstance(ret_hidden_states, tuple): + last_hidden_states, hidden_states = ret_hidden_states + else: last_hidden_states = ret_hidden_states hidden_states = ret_hidden_states - else: - last_hidden_states, hidden_states = ret_hidden_states return last_hidden_states, hidden_states + def _validate_local_argmax_reduction(self) -> None: + if not self.use_local_argmax_reduction: + return + if self.speculative_config.draft_sample_method == "probabilistic": + raise ValueError( + "use_local_argmax_reduction is not compatible with " + "draft_sample_method='probabilistic'." + ) + if not hasattr(self.model, "get_top_tokens"): + raise ValueError( + "use_local_argmax_reduction is enabled but draft model " + f"{self.model.__class__.__name__} does not implement " + "get_top_tokens()." + ) + logger.info( + "Using local argmax reduction for draft token generation " + "(communication: O(2*tp_size) vs O(vocab_size))." + ) + def _sample_draft( self, - logits: torch.Tensor, + hidden_states: torch.Tensor, idx_mapping: torch.Tensor, pos: torch.Tensor, draft_step: torch.Tensor, draft_logits: torch.Tensor | None, ) -> torch.Tensor: if draft_logits is not None: + logits = self.model.compute_logits(hidden_states) # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise # used for draft and target sampling. return gumbel_sample( @@ -251,8 +289,26 @@ def _sample_draft( output_processed_logits_col=draft_step, use_fp64=self.use_fp64_gumbel, ) - else: - return logits.argmax(dim=-1) + if self.use_local_argmax_reduction: + return self.model.get_top_tokens(hidden_states) + logits = self.model.compute_logits(hidden_states) + return logits.argmax(dim=-1) + + def _mtp_prefill_begin(self) -> None: + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(False) + + def _mtp_prefill_end(self, num_reqs: int) -> None: + if self.share_mtp_topk_indices and self.num_speculative_steps > 1: + self.model.model.compact_topk_indices(self.last_token_indices[:num_reqs]) + + def _mtp_decode_begin(self) -> None: + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(True) + + def _mtp_decode_end(self) -> None: + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(False) def prefill( self, @@ -277,10 +333,9 @@ def prefill( mm_inputs=mm_inputs, ) sample_hidden_states = last_hidden_states[last_token_indices] - logits = self.model.compute_logits(sample_hidden_states) self.draft_tokens[:num_reqs, 0] = self._sample_draft( - logits, + sample_hidden_states, idx_mapping, pos, self.current_draft_step, @@ -362,9 +417,8 @@ def generate_draft( last_hidden_states = last_hidden_states[:num_reqs] # Sample the draft tokens. - logits = self.model.compute_logits(last_hidden_states) draft_tokens = self._sample_draft( - logits, + last_hidden_states, idx_mapping, positions, self.current_draft_step, @@ -432,11 +486,13 @@ def capture( # For PIECEWISE, only the model's compiled regions are captured # and the rest (compute_logits, gumbel_sample) runs eagerly. assert self.prefill_cudagraph_manager is not None + self._mtp_prefill_begin() self.prefill_cudagraph_manager.capture( self.prefill, attn_states, progress_bar_desc="Capturing eagle prefill CUDA graphs", ) + self._mtp_prefill_end(self.max_num_reqs) if self.num_speculative_steps == 1: return @@ -445,15 +501,19 @@ def capture( # compute_logits + sample + update_eagle_inputs) for a single # step. assert self.decode_cudagraph_manager is not None - self.decode_cudagraph_manager.capture( - self.generate_draft, - self.model_state, - self.input_buffers, - self.block_tables, - self.attn_groups, - self.kv_cache_config, - progress_bar_desc="Capturing eagle decode CUDA graphs", - ) + self._mtp_decode_begin() + try: + self.decode_cudagraph_manager.capture( + self.generate_draft, + self.model_state, + self.input_buffers, + self.block_tables, + self.attn_groups, + self.kv_cache_config, + progress_bar_desc="Capturing eagle decode CUDA graphs", + ) + finally: + self._mtp_decode_end() @torch.inference_mode() def propose( @@ -548,6 +608,7 @@ def propose( need_eager=is_profile, ) + self._mtp_prefill_begin() if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL: # Replay the full graph for draft prefill. assert self.prefill_cudagraph_manager is not None @@ -565,6 +626,7 @@ def propose( cudagraph_runtime_mode=prefill_batch_desc.cg_mode, mm_inputs=mm_inputs, ) + self._mtp_prefill_end(num_reqs) if self.num_speculative_steps == 1: # Early exit. @@ -593,12 +655,16 @@ def propose( ) # Generate the remaining num_speculative_steps - 1 draft tokens. - self.multi_step_decode( - num_reqs, - dummy_run and skip_attn_for_dummy_run, - decode_batch_desc, - num_tokens_across_dp, - ) + self._mtp_decode_begin() + try: + self.multi_step_decode( + num_reqs, + dummy_run and skip_attn_for_dummy_run, + decode_batch_desc, + num_tokens_across_dp, + ) + finally: + self._mtp_decode_end() return self.draft_tokens[:num_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index e721e85b35..be6fb3d8aa 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -19,10 +19,20 @@ def __init__(self, device: torch.device | None = None): self.num_draft_tokens: int = 0 def set_draft_tokens( - self, input_batch: InputBatch, draft_tokens: torch.Tensor + self, + input_batch: InputBatch, + draft_tokens: torch.Tensor, + num_draft_tokens: int | None = None, ) -> None: self.req_ids = input_batch.req_ids - self.num_draft_tokens = draft_tokens.shape[1] + if num_draft_tokens is None: + num_draft_tokens = draft_tokens.shape[1] + if not 0 <= num_draft_tokens <= draft_tokens.shape[1]: + raise ValueError( + "num_draft_tokens must be within the proposal tensor width" + ) + self.num_draft_tokens = num_draft_tokens + draft_tokens = draft_tokens[:, :num_draft_tokens] if not input_batch.has_structured_output_reqs: # No draft token validation needs to be performed by # the scheduler for this batch. diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 1cdcf567d6..9bcafe5210 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -15,7 +15,13 @@ NewRequestData, SchedulerOutput, ) -from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec, MambaSpec +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + CrossAttentionSpec, + KVCacheSpec, + MambaSpec, + UniformTypeKVCacheSpecs, +) from vllm.v1.request import Request from vllm.v1.worker.gpu.model_runner import GPUModelRunner @@ -29,6 +35,12 @@ def _reserved_block_count( max_encoder_len: int, ) -> int: """Match the scheduler's block reservation in hand-built warmup batches.""" + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): + specs = tuple(kv_cache_spec.kv_cache_specs.values()) + assert specs + kv_cache_spec = specs[0] + if isinstance(kv_cache_spec, CircularBufferSpec): + return 1 if isinstance(kv_cache_spec, CrossAttentionSpec): return cdiv(max_encoder_len, kv_cache_spec.block_size) num_speculative_blocks = 0 diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index c150cf791e..56a4e69e7e 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -153,6 +153,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, + CircularBufferSpec, CrossAttentionSpec, EncoderOnlyAttentionSpec, FullAttentionSpec, @@ -218,6 +219,7 @@ update_ngram_gpu_tensors_incremental, update_scheduler_for_invalid_drafts, ) +from vllm.v1.spec_decode.qwen4_exp import Qwen4ExpMTPProposer from vllm.v1.spec_decode.static_draft_vocab import ( DynamicDraftVocabPrefillBootstrapState, resolve_mtp_draft_vocab_config, @@ -254,6 +256,7 @@ KVBlockZeroer, add_kv_sharing_layers_to_kv_cache_groups, bind_kv_cache, + compressed_kernel_block_size, prepare_kernel_block_sizes, sanity_check_mm_encoder_outputs, ) @@ -688,6 +691,18 @@ def _sync_sm70_before_compile_graph_capture( torch.accelerator.synchronize() +def _select_dummy_sample_hidden_states( + hidden_states: torch.Tensor | IntermediateTensors, + num_scheduled_tokens: np.ndarray, + device: torch.device, +) -> torch.Tensor | None: + if isinstance(hidden_states, IntermediateTensors): + return None + logit_indices = np.cumsum(num_scheduled_tokens) - 1 + logit_indices_device = torch.from_numpy(logit_indices).to(device, non_blocking=True) + return hidden_states[logit_indices_device] + + def _sm70_profile_trace(message: str, *args: object) -> None: if envs.VLLM_SM70_PROFILE_TRACE: if args: @@ -1380,6 +1395,19 @@ def __init__( # Only relevant for models using ALiBi (e.g, MPT) self.use_alibi = model_config.uses_alibi + ple_layer_ids = getattr(model_config.hf_text_config, "ple_layer_ids", ()) + self.uses_ngram_embedding = bool(ple_layer_ids) + if self.uses_ngram_embedding: + self.ngram_context_len = int(model_config.hf_text_config.ngram_size) - 1 + self.ngram_eos_token_id = int(model_config.hf_text_config.eos_token_id) + else: + self.ngram_context_len = 0 + self.ngram_eos_token_id = 0 + if self.uses_ngram_embedding and self.ngram_context_len <= 0: + raise ValueError("N-gram embedding requires context length >= 1") + if self.uses_ngram_embedding and parallel_config.pipeline_parallel_size > 1: + raise RuntimeError("N-gram PLE embedding requires pipeline_parallel_size=1") + self.cascade_attn_enabled = not self.model_config.disable_cascade_attn self.is_mm_prefix_lm = self.model_config.is_mm_prefix_lm @@ -1504,6 +1532,7 @@ def __init__( | ExtractHiddenStatesProposer | Gemma4Proposer | Step3p5MTPProposer + | Qwen4ExpMTPProposer ) if self.speculative_config.method == "custom_class": self.drafter = create_custom_proposer( # type: ignore[assignment] @@ -1540,6 +1569,8 @@ def __init__( self.drafter = Gemma4Proposer(self.vllm_config, self.device, self) elif self.speculative_config.use_step3p5_mtp(): self.drafter = Step3p5MTPProposer(self.vllm_config, self.device, self) + elif self.speculative_config.use_qwen4_exp_mtp(): + self.drafter = Qwen4ExpMTPProposer(self.vllm_config, self.device, self) elif self.speculative_config.use_dspark(): self.drafter = DSparkProposer(self.vllm_config, self.device, self) self.use_aux_hidden_state_outputs = True @@ -1713,6 +1744,12 @@ def __init__( self.inputs_embeds = self._make_buffer( self.max_num_tokens, self.inputs_embeds_size, dtype=self.dtype, numpy=False ) + if self.uses_ngram_embedding: + self.ngram_context = self._make_buffer( + self.max_num_reqs, + self.ngram_context_len, + dtype=torch.int32, + ) self.is_token_ids = self._make_buffer(self.max_num_tokens, dtype=torch.bool) self.discard_request_mask = self._make_buffer( self.max_num_reqs, dtype=torch.bool @@ -5088,11 +5125,26 @@ def _prepare_inputs( profile_positions_ms = (time.perf_counter() - profile_inner_t0) * 1000.0 profile_inner_t0 = time.perf_counter() if profile_inputs else 0.0 - self.input_batch.block_table.compute_slot_mapping( - num_reqs, - self.query_start_loc.gpu[: num_reqs + 1], - self.positions[:total_num_scheduled_tokens], - ) + for group_id, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups): + kv_cache_spec = kv_cache_group.kv_cache_spec + is_circular = isinstance(kv_cache_spec, CircularBufferSpec) + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): + is_circular = all( + isinstance(spec, CircularBufferSpec) + for spec in kv_cache_spec.kv_cache_specs.values() + ) + block_table = self.input_batch.block_table[group_id] + if is_circular: + # QSA derives ring slots from logical positions in its own + # metadata builder. Generic mapping would index beyond the + # ring's single block-table column. + block_table.slot_mapping.gpu.fill_(PAD_SLOT_ID) + continue + block_table.compute_slot_mapping( + num_reqs, + self.query_start_loc.gpu[: num_reqs + 1], + self.positions[:total_num_scheduled_tokens], + ) if profile_inputs: profile_slot_mapping_inner_ms = ( time.perf_counter() - profile_inner_t0 @@ -5629,7 +5681,15 @@ def _build_attn_group_metadata( else: spec_decode_common_attn_metadata = cm # Capture per-group block tables for multi-group proposers. - if self.speculative_config and isinstance(self.drafter, Step3p5MTPProposer): + if self.speculative_config and isinstance( + self.drafter, Qwen4ExpMTPProposer + ): + self.drafter.set_per_group_block_table( + kv_cache_gid, cm.block_table_tensor + ) + elif self.speculative_config and isinstance( + self.drafter, Step3p5MTPProposer + ): self.drafter.set_per_group_attn_metadata( kv_cache_gid, cm.block_table_tensor, cm.slot_mapping ) @@ -6593,10 +6653,83 @@ def _prepare_mm_inputs( inputs_embeds = self.inputs_embeds.gpu[:num_tokens] return input_ids, inputs_embeds + def _prepare_ngram_context( + self, + num_reqs: int, + num_reqs_padded: int, + ) -> torch.Tensor: + """Copy committed per-request token history into the PLE context.""" + if not self.uses_ngram_embedding: + raise RuntimeError("N-gram context requested for a non-PLE model") + + eos_token_id = int(self.ngram_eos_token_id) + context_cpu = self.ngram_context.np[:num_reqs_padded] + context_cpu.fill(eos_token_id) + num_computed = self.input_batch.num_computed_tokens_cpu + token_ids = self.input_batch.token_ids_cpu + is_token_ids = self.input_batch.is_token_ids + + for req_idx in range(num_reqs): + end = int(num_computed[req_idx]) + if end <= 0: + continue + start = max(0, end - self.ngram_context_len) + context_tokens = token_ids[req_idx, start:end] + if context_tokens.size == 0: + continue + if self.enable_prompt_embeds and not is_token_ids[req_idx, start:end].all(): + context_tokens = context_tokens.copy() + context_tokens[~is_token_ids[req_idx, start:end]] = eos_token_id + context_cpu[req_idx, -context_tokens.size :] = context_tokens + + self._copy_buffer_to_gpu(self.ngram_context, num_reqs_padded) + return self.ngram_context.gpu[:num_reqs_padded] + + def _maybe_add_ngram_kwargs( + self, + model_kwargs: dict[str, Any], + *, + num_reqs: int, + num_reqs_padded: int, + is_first_rank: bool, + is_encoder_decoder: bool, + use_dummy_context: bool, + query_start_loc: torch.Tensor | None = None, + num_scheduled_tokens: Sequence[int] | np.ndarray | None = None, + ) -> None: + if not self.uses_ngram_embedding or not is_first_rank or is_encoder_decoder: + return + + if query_start_loc is None: + if num_scheduled_tokens is None: + raise RuntimeError("query_start_loc is required for N-gram PLE") + scheduled = np.asarray(num_scheduled_tokens, dtype=np.int32) + cu_num_tokens = np.cumsum(scheduled, dtype=np.int32) + last = int(cu_num_tokens[-1]) if num_reqs > 0 else 0 + self.query_start_loc.np[0] = 0 + if num_reqs > 0: + self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens + self.query_start_loc.np[num_reqs + 1 :].fill(last) + self._copy_buffer_to_gpu(self.query_start_loc) + query_start_loc = self.query_start_loc.gpu[: num_reqs_padded + 1] + model_kwargs["query_start_loc"] = query_start_loc + + if use_dummy_context: + self.ngram_context.np[:num_reqs_padded].fill(int(self.ngram_eos_token_id)) + self._copy_buffer_to_gpu(self.ngram_context, num_reqs_padded) + model_kwargs["ngram_context"] = self.ngram_context.gpu[:num_reqs_padded] + else: + model_kwargs["ngram_context"] = self._prepare_ngram_context( + num_reqs, + num_reqs_padded, + ) + def _preprocess( self, scheduler_output: "SchedulerOutput", num_input_tokens: int, # Padded + num_reqs: int, + num_reqs_padded: int, intermediate_tensors: IntermediateTensors | None = None, ) -> tuple[ torch.Tensor | None, @@ -6675,6 +6808,25 @@ def _preprocess( inputs_embeds = None model_kwargs = self._init_model_kwargs() + if ( + self.uses_ngram_embedding + and is_first_rank + and not is_encoder_decoder + and input_ids is None + ): + raise RuntimeError( + "N-gram PLE requires token IDs on the first pipeline rank" + ) + self._maybe_add_ngram_kwargs( + model_kwargs, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + is_first_rank=is_first_rank, + is_encoder_decoder=is_encoder_decoder, + use_dummy_context=False, + query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], + ) + if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_input_tokens] elif self.uses_xdrope_dim > 0: @@ -8540,7 +8692,11 @@ def execute_model( model_kwargs, ec_connector_output, ) = self._preprocess( - scheduler_output, num_tokens_padded, intermediate_tensors + scheduler_output, + num_tokens_padded, + num_reqs, + num_reqs_padded, + intermediate_tensors, ) if trace_log: trace_model_preprocess_ms = ( @@ -10496,7 +10652,10 @@ def _dummy_run( num_active_loras: int = 0, profile_seq_lens: int | None = None, batch_descriptor_override: BatchDescriptor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[ + torch.Tensor | IntermediateTensors, + torch.Tensor | None, + ]: """ Run a dummy forward pass to warm up/profile run or capture the CUDA graph for the model. @@ -10832,6 +10991,16 @@ def _dummy_run( input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None + self._maybe_add_ngram_kwargs( + model_kwargs, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + is_first_rank=get_pp_group().is_first_rank, + is_encoder_decoder=self.model_config.is_encoder_decoder, + use_dummy_context=True, + num_scheduled_tokens=num_scheduled_tokens, + ) + if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] elif self.uses_xdrope_dim > 0: @@ -10984,16 +11153,19 @@ def _dummy_run( self.eplb_step(is_dummy=True, is_profile=is_profile) _sm70_profile_trace("_dummy_run eplb exit") - logit_indices = np.cumsum(num_scheduled_tokens) - 1 - logit_indices_device = torch.from_numpy(logit_indices).to( - self.device, non_blocking=True + sample_hidden_states = _select_dummy_sample_hidden_states( + hidden_states, num_scheduled_tokens, self.device ) _sm70_profile_trace( - "_dummy_run return hidden_shape=%s sampled_count=%s", - tuple(hidden_states.shape), - len(logit_indices), + "_dummy_run return hidden_type=%s sampled_shape=%s", + type(hidden_states).__name__, + ( + tuple(sample_hidden_states.shape) + if sample_hidden_states is not None + else None + ), ) - return hidden_states, hidden_states[logit_indices_device] + return hidden_states, sample_hidden_states @torch.inference_mode() def _dummy_sampler_run( @@ -11277,11 +11449,17 @@ def profile_run(self) -> None: self.max_num_tokens, is_profile=True ) _sm70_profile_trace( - "profile_run dummy_run exit hidden_shape=%s last_hidden_shape=%s", - tuple(hidden_states.shape), - tuple(last_hidden_states.shape), + "profile_run dummy_run exit hidden_type=%s last_hidden_shape=%s", + type(hidden_states).__name__, + ( + tuple(last_hidden_states.shape) + if last_hidden_states is not None + else None + ), ) if get_pp_group().is_last_rank: + assert isinstance(hidden_states, torch.Tensor) + assert last_hidden_states is not None if self.is_pooling_model: _sm70_profile_trace("profile_run pooler enter") output = self._dummy_pooler_run(hidden_states) @@ -12049,9 +12227,12 @@ def may_reinitialize_input_batch( continue block_size = kv_cache_group.kv_cache_spec.block_size block_sizes.append(block_size) - max_num_blocks_per_req = cdiv( - max_model_len, block_size * get_total_cp_world_size() - ) + if isinstance(kv_cache_group.kv_cache_spec, CircularBufferSpec): + max_num_blocks_per_req = 1 + else: + max_num_blocks_per_req = cdiv( + max_model_len, block_size * get_total_cp_world_size() + ) if isinstance(kv_cache_group.kv_cache_spec, MambaSpec): max_num_blocks_per_req = ( max_num_blocks_per_req @@ -12166,15 +12347,18 @@ def _reshape_kv_cache_tensors( num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes if isinstance(kv_cache_spec, AttentionSpec): has_attn = True - num_blocks_per_kv_block = ( - kv_cache_spec.block_size // kernel_block_size - ) - kernel_num_blocks = num_blocks * num_blocks_per_kv_block - - # For MLA with compression, storage_block_size != block_size + # Compressed MLA is split into physical pool pages instead + # of the scheduler's attention-kernel blocks. if kv_cache_spec.storage_block_size != kv_cache_spec.block_size: - shape_block_size = kv_cache_spec.storage_block_size + shape_block_size = compressed_kernel_block_size(kv_cache_spec) + kernel_num_blocks = num_blocks * ( + kv_cache_spec.storage_block_size // shape_block_size + ) else: + num_blocks_per_kv_block = ( + kv_cache_spec.block_size // kernel_block_size + ) + kernel_num_blocks = num_blocks * num_blocks_per_kv_block shape_block_size = kernel_block_size kv_cache_shape = attn_backend.get_kv_cache_shape( diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 7ffb0c28f9..ef363d610a 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -13,12 +13,14 @@ from vllm.config import CacheConfig from vllm.logger import init_logger from vllm.model_executor.layers.mamba.mamba_utils import ( - MambaStateCopyFunc, + MambaStateCopyFuncs, + MambaStateCopyFuncsByType, get_conv_copy_spec, get_temporal_copy_spec, ) from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.utils import CpuGpuBuffer @@ -77,8 +79,7 @@ def postprocess_mamba_fused_kernel( # block table. block_table_ptrs_ptr, block_table_stride_req: tl.int64, # stride between requests (in elements) - # Mamba state metadata (per-layer, per-state-type) - # These are 1D arrays indexed by (layer_idx * num_state_types + state_type_idx) + # Mamba state metadata, flattened across all layers and state types. state_base_addrs_ptr, # base address of each state tensor state_block_strides_ptr, # bytes per block for each state state_elem_sizes_ptr, # element size for each state @@ -101,13 +102,13 @@ def postprocess_mamba_fused_kernel( Fused GPU kernel for postprocess_mamba that computes decisions AND performs mamba state copies without any CPU-GPU synchronization. - Grid: (num_reqs, num_layers * num_state_types) + Grid: (num_reqs, num_states) - program_id(0) = request index - program_id(1) = state_idx (flattened index into layer/state_type metadata) - Note: num_layers and num_state_types are not passed as kernel parameters - because the kernel indexes directly into pre-flattened metadata arrays - using program_id(1). The grid dimensions encode the total state count. + The kernel indexes directly into pre-flattened metadata arrays using + program_id(1). The grid dimensions encode the total state count, allowing + different Mamba layer types to own different numbers of state tensors. """ req_idx = tl.program_id(0) state_idx = tl.program_id(1) @@ -307,6 +308,9 @@ def warmup_batch_memcpy_kernel(device: torch.device) -> bool: return True +MambaStateCopyFuncsInput = MambaStateCopyFuncs | MambaStateCopyFuncsByType + + def get_mamba_groups(kv_cache_config: KVCacheConfig) -> tuple[list[int], MambaSpec]: mamba_group_ids: list[int] = [] mamba_specs: list[MambaSpec] = [] @@ -316,10 +320,52 @@ def get_mamba_groups(kv_cache_config: KVCacheConfig) -> tuple[list[int], MambaSp mamba_group_ids.append(i) mamba_specs.append(kv_cache_spec) assert len(mamba_group_ids) > 0, "no mamba layers in the model" - assert all(mamba_specs[0] == spec for spec in mamba_specs) + # State shapes, dtypes, backend types, and TP replication may differ across + # groups (for example Qwen4Exp GDN versus its PLE short-conv state). The + # align scheduler and copy kernels only require shared block semantics. + for field in ("block_size", "mamba_cache_mode", "num_speculative_blocks"): + values = {getattr(spec, field) for spec in mamba_specs} + assert len(values) == 1, f"all mamba groups must share {field}, got {values}" return mamba_group_ids, mamba_specs[0] +def get_mamba_types( + kv_cache_config: KVCacheConfig, +) -> set[MambaAttentionBackendEnum]: + mamba_group_ids, _ = get_mamba_groups(kv_cache_config) + mamba_types: set[MambaAttentionBackendEnum] = set() + for group_id in mamba_group_ids: + spec = kv_cache_config.kv_cache_groups[group_id].kv_cache_spec + assert isinstance(spec, MambaSpec) + mamba_types.add(spec.mamba_type) + return mamba_types + + +def _get_copy_funcs_for_group( + kv_cache_config: KVCacheConfig, + mamba_group_id: int, + copy_funcs: MambaStateCopyFuncsInput, +) -> MambaStateCopyFuncs: + spec = kv_cache_config.kv_cache_groups[mamba_group_id].kv_cache_spec + assert isinstance(spec, MambaSpec) + if isinstance(copy_funcs, dict): + group_copy_funcs = copy_funcs[spec.mamba_type] + assert len(group_copy_funcs) == len(spec.shapes), ( + f"{spec.mamba_type} exposes {len(spec.shapes)} states but has " + f"{len(group_copy_funcs)} copy functions" + ) + return group_copy_funcs + + # Backward compatibility for homogeneous models and legacy callers. Some + # mixed models append a one-state short-conv group after a two-state GDN + # group, so selecting the matching prefix is safe for their old tuple. + assert len(copy_funcs) >= len(spec.shapes), ( + f"{spec.mamba_type} exposes {len(spec.shapes)} states but has only " + f"{len(copy_funcs)} copy functions" + ) + return copy_funcs[: len(spec.shapes)] + + @dataclasses.dataclass class MambaCopyBuffers: src_ptrs: CpuGpuBuffer @@ -334,15 +380,16 @@ def create( cls, max_num_reqs: int, kv_cache_config: KVCacheConfig, - copy_funcs: tuple[MambaStateCopyFunc, ...], + copy_funcs: MambaStateCopyFuncsInput, make_buffer: Callable[..., CpuGpuBuffer], copies_per_req: int = 1, ) -> "MambaCopyBuffers": mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config) entries_per_req = sum( len(kv_cache_config.kv_cache_groups[gid].layer_names) + * len(_get_copy_funcs_for_group(kv_cache_config, gid, copy_funcs)) for gid in mamba_group_ids - ) * len(copy_funcs) + ) n = max_num_reqs * entries_per_req * copies_per_req return cls( src_ptrs=make_buffer(n, dtype=torch.int64), @@ -369,7 +416,7 @@ class MambaSpecDecodeGPUContext: window with offset-based copies), 0 for temporal states (full block copies). """ - # Per-state metadata tensors (shape: [num_layers * num_state_types]) + # Per-state metadata tensors (shape: [num_states]) # These are populated from forward_context during the first forward pass state_base_addrs: torch.Tensor # int64: base address of each state tensor state_block_strides: torch.Tensor # int64: bytes per block @@ -381,7 +428,7 @@ class MambaSpecDecodeGPUContext: # Configuration block_size: int num_layers: int - num_state_types: int + num_states: int mamba_group_ids: list[int] num_groups: int @@ -412,7 +459,7 @@ def create( cls, max_num_reqs: int, kv_cache_config: KVCacheConfig, - num_state_types: int, + mamba_state_copy_funcs: MambaStateCopyFuncsInput, device: torch.device, make_buffer: Callable[..., CpuGpuBuffer], ) -> "MambaSpecDecodeGPUContext": @@ -424,30 +471,28 @@ def create( len(kv_cache_config.kv_cache_groups[gid].layer_names) for gid in mamba_group_ids ) - total_states = num_layers * num_state_types + num_states = sum( + len(kv_cache_config.kv_cache_groups[gid].layer_names) + * len( + _get_copy_funcs_for_group(kv_cache_config, gid, mamba_state_copy_funcs) + ) + for gid in mamba_group_ids + ) return cls( - state_base_addrs=torch.zeros( - total_states, dtype=torch.int64, device=device - ), + state_base_addrs=torch.zeros(num_states, dtype=torch.int64, device=device), state_block_strides=torch.zeros( - total_states, dtype=torch.int64, device=device - ), - state_elem_sizes=torch.zeros( - total_states, dtype=torch.int32, device=device - ), - state_inner_sizes=torch.zeros( - total_states, dtype=torch.int64, device=device - ), - state_conv_widths=torch.zeros( - total_states, dtype=torch.int32, device=device + num_states, dtype=torch.int64, device=device ), + state_elem_sizes=torch.zeros(num_states, dtype=torch.int32, device=device), + state_inner_sizes=torch.zeros(num_states, dtype=torch.int64, device=device), + state_conv_widths=torch.zeros(num_states, dtype=torch.int32, device=device), state_group_indices=torch.zeros( - total_states, dtype=torch.int32, device=device + num_states, dtype=torch.int32, device=device ), block_size=mamba_spec.block_size, num_layers=num_layers, - num_state_types=num_state_types, + num_states=num_states, mamba_group_ids=mamba_group_ids, num_groups=len(mamba_group_ids), num_accepted_tokens_out=torch.zeros( @@ -470,7 +515,7 @@ def initialize_from_forward_context( self, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + mamba_state_copy_funcs: MambaStateCopyFuncsInput, block_tables: list[torch.Tensor], ) -> None: """ @@ -501,8 +546,9 @@ def initialize_from_forward_context( forward_context: Dictionary mapping layer names to attention objects, populated after the model is loaded. Each attention object must have a `kv_cache` attribute containing the list of state tensors. - mamba_state_copy_funcs: Tuple of copy functions (one per state type) - used to determine whether each state is a conv or temporal state. + mamba_state_copy_funcs: Copy functions keyed by Mamba backend type, + or a legacy homogeneous tuple. They determine whether each state + is a conv or temporal state. block_tables: per-mamba-group persistent block-table tensors, in the same order as `mamba_group_ids`. Their `data_ptr()` / `stride(0)` are captured once for the kernel to index into. @@ -513,11 +559,18 @@ def initialize_from_forward_context( idx = 0 for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids): layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names + group_copy_funcs = _get_copy_funcs_for_group( + kv_cache_config, mamba_group_id, mamba_state_copy_funcs + ) for layer_name in layer_names: attention = forward_context[layer_name] kv_caches: list[torch.Tensor] = attention.kv_cache + assert len(kv_caches) == len(group_copy_funcs), ( + f"{layer_name} has {len(kv_caches)} state tensors but " + f"{len(group_copy_funcs)} copy functions" + ) - for state_type_idx, state in enumerate(kv_caches): + for state, copy_func in zip(kv_caches, group_copy_funcs): # Base address self.state_base_addrs[idx] = state.data_ptr() @@ -534,7 +587,6 @@ def initialize_from_forward_context( # Element size self.state_elem_sizes[idx] = state.element_size() - copy_func = mamba_state_copy_funcs[state_type_idx] assert ( copy_func is get_conv_copy_spec or copy_func is get_temporal_copy_spec @@ -566,6 +618,10 @@ def initialize_from_forward_context( self.state_group_indices[idx] = group_local_idx idx += 1 + assert idx == self.num_states, ( + f"initialized {idx} mamba states, expected {self.num_states}" + ) + # Cache per-group block-table base addresses and per-request stride. # `block_tables[i]` is the persistent 2D int32 block-table tensor for # `mamba_group_ids[i]`; `data_ptr()` / `stride(0)` are stable for the @@ -586,7 +642,7 @@ def initialize_from_forward_context( "postprocess_init", num_groups=self.num_groups, num_layers=self.num_layers, - num_state_types=self.num_state_types, + num_states=self.num_states, block_size=self.block_size, block_table_stride_req=self.block_table_stride_req, ) @@ -628,8 +684,7 @@ def run_fused_postprocess( spec_state_slot_selectors_gpu[:num_reqs] ) - total_states = self.num_layers * self.num_state_types - grid = (num_reqs, total_states) + grid = (num_reqs, self.num_states) has_ddtree = ddtree_accepted_node_indices is not None if ddtree_accepted_node_indices is None: ddtree_accepted_node_indices = num_accepted_tokens_gpu.new_empty((1, 1)) @@ -665,7 +720,6 @@ def warmup_fused_postprocess(self) -> bool: return False device = self.state_base_addrs.device - total_states = self.num_layers * self.num_state_types warmup_sizes = {1} if envs.VLLM_SM70_MTP_CONCURRENCY_WARMUP: max_num_reqs = int(self.num_accepted_tokens_out.shape[0]) @@ -689,7 +743,7 @@ def warmup_fused_postprocess(self) -> bool: (1, 1), dtype=torch.int32, device=device ) - postprocess_mamba_fused_kernel[(num_reqs, total_states)]( + postprocess_mamba_fused_kernel[(num_reqs, self.num_states)]( num_accepted_tokens, spec_state_slot_selectors, ddtree_accepted_node_indices, @@ -735,7 +789,7 @@ def create( cls, max_num_reqs: int, kv_cache_config: KVCacheConfig, - copy_funcs: tuple[MambaStateCopyFunc, ...], + copy_funcs: MambaStateCopyFuncsInput, make_buffer: Callable[..., CpuGpuBuffer], device: torch.device, with_postprocess_align: bool, @@ -748,7 +802,7 @@ def create( MambaSpecDecodeGPUContext.create( max_num_reqs=max_num_reqs, kv_cache_config=kv_cache_config, - num_state_types=len(copy_funcs), + mamba_state_copy_funcs=copy_funcs, device=device, make_buffer=make_buffer, ) @@ -761,7 +815,7 @@ def create( def collect_mamba_copy_meta( copy_bufs: MambaCopyBuffers, kv_cache_config: KVCacheConfig, - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + mamba_state_copy_funcs: MambaStateCopyFuncsInput, mamba_group_ids: list[int], src_block_idx: int, dest_block_idx: int, @@ -781,13 +835,20 @@ def collect_mamba_copy_meta( offset = copy_bufs.offset for mamba_group_id in mamba_group_ids: + group_copy_funcs = _get_copy_funcs_for_group( + kv_cache_config, mamba_group_id, mamba_state_copy_funcs + ) block_ids = req_state.block_ids[mamba_group_id] dest_block_id = block_ids[dest_block_idx] layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names for layer_name in layer_names: attention = forward_context[layer_name] kv_caches: list[torch.Tensor] = attention.kv_cache - for state, state_copy_func in zip(kv_caches, mamba_state_copy_funcs): + assert len(kv_caches) == len(group_copy_funcs), ( + f"{layer_name} has {len(kv_caches)} state tensors but " + f"{len(group_copy_funcs)} copy functions" + ) + for state, state_copy_func in zip(kv_caches, group_copy_funcs): copy_src_block_idx = src_block_idx copy_num_accepted_tokens = accept_token_bias + 1 if state_slot_bias != accept_token_bias: @@ -844,7 +905,7 @@ def preprocess_mamba( input_batch: GPUInputBatch, requests: dict[str, CachedRequestState], forward_context: dict[str, Any], - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + mamba_state_copy_funcs: MambaStateCopyFuncsInput, copy_bufs: MambaCopyBuffers, ): """ @@ -968,7 +1029,7 @@ def postprocess_mamba( requests: dict[str, CachedRequestState], mamba_state_idx: dict[str, int], forward_context: dict[str, Any], - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + mamba_state_copy_funcs: MambaStateCopyFuncsInput, copy_bufs: MambaCopyBuffers, ddtree_accepted_node_indices: torch.Tensor | None = None, ): @@ -1089,7 +1150,7 @@ def postprocess_mamba_align_gpu( input_batch: GPUInputBatch, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + mamba_state_copy_funcs: MambaStateCopyFuncsInput, ddtree_accepted_node_indices: torch.Tensor | None = None, ) -> None: """GPU-side mamba postprocess for spec decode + hybrid + align mode. @@ -1123,7 +1184,7 @@ def postprocess_mamba_align_gpu( "postprocess_align_gpu", num_reqs=num_reqs, num_layers=ctx.num_layers, - num_state_types=ctx.num_state_types, + num_states=ctx.num_states, block_size=ctx.block_size, ) ctx.run_fused_postprocess( diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 6b634557a5..e6da3074c2 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -16,7 +16,7 @@ from vllm.model_executor.models.utils import extract_layer_index from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.math_utils import largest_power_of_2_divisor +from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.v1.attention.backend import ( AttentionBackend, @@ -33,10 +33,20 @@ MambaSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.worker.block_table import get_block_table_width logger = init_logger(__name__) +def compressed_kernel_block_size(spec: AttentionSpec) -> int: + storage = spec.storage_block_size + max_page = max(PAGED_MQA_PAGE_SIZES) + min_page = min(PAGED_MQA_PAGE_SIZES) + if storage <= max_page: + return storage + return max_page if storage % max_page == 0 else min_page + + def _infer_segment_block_strides( seg_addrs: list[int], page_size_el: int, @@ -75,14 +85,12 @@ def _infer_segment_block_strides( return block_strides -@triton.jit(do_not_specialize=["n_blocks"]) +@triton.jit def _zero_kv_blocks_kernel( seg_addrs_ptr, seg_block_strides_ptr, + seg_page_sizes_ptr, block_ids_ptr, - n_blocks, - N_SEGS: tl.constexpr, - PAGE_SIZE_EL: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): """Zero KV cache blocks across all segments in a single launch. @@ -98,27 +106,27 @@ def _zero_kv_blocks_kernel( elements, kept separate from the PAGE_SIZE_EL span each block zeros (they differ for interleaved layouts, see init_meta). - Programs are mapped as (block_index, seg_index, chunk_index). + Programs are mapped as (block_index, seg_index, chunk_index). Segment page + sizes may differ across hybrid cache groups. """ - pid = tl.program_id(0) - chunks = PAGE_SIZE_EL // BLOCK_SIZE - work_per_block = N_SEGS * chunks - block_index = pid // work_per_block - if block_index >= n_blocks: + block_index = tl.program_id(0) + seg_index = tl.program_id(1) + chunk_index = tl.program_id(2) + page_size_el = tl.load(seg_page_sizes_ptr + seg_index) + chunk_offset = chunk_index.to(tl.int64) * BLOCK_SIZE + if chunk_offset >= page_size_el: return - remainder = pid % work_per_block - seg_index = remainder // chunks - chunk_index = remainder % chunks block_id = tl.load(block_ids_ptr + block_index) seg_addr = tl.load(seg_addrs_ptr + seg_index) block_stride_el = tl.load(seg_block_strides_ptr + seg_index) ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32)) - offset = ( - block_id.to(tl.int64) * block_stride_el.to(tl.int64) - + chunk_index.to(tl.int64) * BLOCK_SIZE + block_offset = block_id.to(tl.int64) * block_stride_el.to(tl.int64) + cols = chunk_offset + tl.arange(0, BLOCK_SIZE).to(tl.int64) + tl.store( + ptr + block_offset + cols, + tl.zeros([BLOCK_SIZE], dtype=tl.int32), + mask=cols < page_size_el, ) - cols = tl.arange(0, BLOCK_SIZE).to(tl.int64) - tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32)) class KVBlockZeroer: @@ -132,7 +140,9 @@ class KVBlockZeroer: def __init__(self, device: torch.device, pin_memory: bool): self.device = device self.pin_memory = pin_memory - self._meta: tuple[torch.Tensor, torch.Tensor, int, int, int] | None = None + self._meta: ( + tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int] | None + ) = None self._id_cap: int = 0 self._ids_pinned: torch.Tensor | None = None self._ids_gpu: torch.Tensor | None = None @@ -161,9 +171,17 @@ def init_meta( state tensor views over one raw page per block; the first state tensor starts at the page base, so one segment per layer zeros the whole page. """ - seen_ptrs: set[int] = set() + seen_ptrs: dict[int, int] = {} seg_addrs: list[int] = [] - page_size_el: int | None = None + seg_page_sizes: list[int] = [] + + def add_segment(address: int, page_size_el: int) -> None: + if (index := seen_ptrs.get(address)) is not None: + seg_page_sizes[index] = max(seg_page_sizes[index], page_size_el) + return + seen_ptrs[address] = len(seg_addrs) + seg_addrs.append(address) + seg_page_sizes.append(page_size_el) for group in attn_groups_iter: spec = group.kv_cache_spec @@ -186,21 +204,12 @@ def init_meta( cache_dtype_str=cache_dtype, ) dp = kv.data_ptr() - if dp in seen_ptrs: - continue - seen_ptrs.add(dp) el = kv.element_size() cur_bytes = kv.stride(block_dim) * el assert cur_bytes % 4 == 0 kernel_block_el = cur_bytes // 4 cur_page_el = kernel_block_el * ratio - if page_size_el is None: - page_size_el = cur_page_el - else: - assert page_size_el == cur_page_el, ( - f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}" - ) block_stride_bytes = cur_bytes outer_dims = [ @@ -211,7 +220,7 @@ def init_meta( outer_strides = [kv.stride(d) * el for d in outer_dims] for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)): off_bytes = sum(i * s for i, s in zip(outer, outer_strides)) - seg_addrs.append(dp + off_bytes) + add_segment(dp + off_bytes, cur_page_el) elif isinstance(spec, MambaSpec): if not isinstance(kv, (list, tuple)) or not kv: continue @@ -219,34 +228,38 @@ def init_meta( if not isinstance(first_state, torch.Tensor): continue dp = first_state.data_ptr() - if dp in seen_ptrs: - continue - seen_ptrs.add(dp) page_bytes = spec.page_size_bytes assert page_bytes % 4 == 0 cur_page_el = page_bytes // 4 - if page_size_el is None: - page_size_el = cur_page_el - else: - assert page_size_el == cur_page_el, ( - f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}" - ) - seg_addrs.append(dp) - - if not seg_addrs or page_size_el is None: + add_segment(dp, cur_page_el) + + if not seg_addrs: self._meta = None return - blk_size = min(largest_power_of_2_divisor(page_size_el), 1024) + max_page_size_el = max(seg_page_sizes) + blk_size = min(1 << (max_page_size_el - 1).bit_length(), 1024) # Dense layouts space blocks page_size_el apart. Block-major # interleaved pools expose segment starts one cell apart and advance by # the number of segments in that contiguous address run. Infer each # run separately because a process may own more than one KV pool. n_segs = len(seg_addrs) + block_strides = [0] * n_segs + for page_size_el in set(seg_page_sizes): + indices = [ + index + for index, size in enumerate(seg_page_sizes) + if size == page_size_el + ] + inferred = _infer_segment_block_strides( + [seg_addrs[index] for index in indices], page_size_el + ) + for index, stride in zip(indices, inferred): + block_strides[index] = stride seg_block_strides = torch.tensor( - _infer_segment_block_strides(seg_addrs, page_size_el), + block_strides, dtype=torch.int64, device=self.device, ) @@ -261,7 +274,8 @@ def init_meta( self._meta = ( torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device), seg_block_strides, - page_size_el, + torch.tensor(seg_page_sizes, dtype=torch.int64, device=self.device), + (max_page_size_el + blk_size - 1) // blk_size, blk_size, n_segs, ) @@ -270,7 +284,14 @@ def zero_block_ids(self, block_ids: list[int]) -> None: """Zero the KV cache memory for the given block IDs.""" if not block_ids or self._meta is None: return - seg_addrs, seg_block_strides, page_size_el, blk_size, n_segs = self._meta + ( + seg_addrs, + seg_block_strides, + seg_page_sizes, + max_chunks, + blk_size, + n_segs, + ) = self._meta n_blocks = len(block_ids) if n_blocks > self._id_cap: self._id_cap = n_blocks * 2 @@ -286,14 +307,12 @@ def zero_block_ids(self, block_ids: list[int]) -> None: self._ids_pinned[:n_blocks].numpy()[:] = block_ids idx = self._ids_gpu[:n_blocks] idx.copy_(self._ids_pinned[:n_blocks], non_blocking=True) - grid = (n_blocks * n_segs * (page_size_el // blk_size),) + grid = (n_blocks, n_segs, max_chunks) _zero_kv_blocks_kernel[grid]( seg_addrs, seg_block_strides, + seg_page_sizes, idx, - n_blocks, - N_SEGS=n_segs, - PAGE_SIZE_EL=page_size_el, BLOCK_SIZE=blk_size, ) @@ -325,20 +344,44 @@ def create_metadata_builders( kernel_block_size: int | None = None, num_metadata_builders: int = 1, ): - kv_cache_spec_builder = ( - self.kv_cache_spec.copy_with_new_block_size(kernel_block_size) - if kernel_block_size is not None - else self.kv_cache_spec - ) + if kernel_block_size is None: + kv_cache_spec_builder = self.kv_cache_spec + elif ( + isinstance(self.kv_cache_spec, AttentionSpec) + and self.kv_cache_spec.storage_block_size != self.kv_cache_spec.block_size + ): + compress_ratio = ( + self.kv_cache_spec.block_size // self.kv_cache_spec.storage_block_size + ) + kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size( + compressed_kernel_block_size(self.kv_cache_spec) * compress_ratio + ) + else: + kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size( + kernel_block_size + ) + builder_cls = self.backend.get_builder_cls() + builder_kwargs = {} + if builder_cls.requires_block_table_width: + max_num_blocks = self.kv_cache_spec.max_num_blocks_per_req( + vllm_config, vllm_config.model_config.max_model_len + ) + builder_kwargs["block_table_width"] = get_block_table_width( + max_num_blocks, self.kv_cache_spec.block_size, kernel_block_size + ) self.metadata_builders = [ - self.backend.get_builder_cls()( + builder_cls( kv_cache_spec_builder, self.layer_names, vllm_config, device, + **builder_kwargs, ) for _ in range(num_metadata_builders) ] + if kernel_block_size is not None: + for builder in self.metadata_builders: + builder.kernel_block_size = kernel_block_size # type: ignore[attr-defined] def get_metadata_builder(self, ubatch_id: int = 0) -> AttentionMetadataBuilder: assert len(self.metadata_builders) > ubatch_id @@ -601,7 +644,7 @@ def bind_kv_cache( # Bind kv_caches to forward context for layer_name, kv_cache in kv_caches.items(): - forward_context[layer_name].kv_cache = kv_cache + forward_context[layer_name].bind_kv_cache(kv_cache) def is_residual_scattered_for_sp(